LLM Speedrun: Architecture

AI Use #

Audience #

Python’s numpyor with

Rust’s rayoncrate (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.

Table of Contents

A note on notation:

• All scalers are represented as matrix.
• All vectors are represented as column vectors, like .
• A matrix of vectors is represented as concatenated column vectors, like .
• represents a matrix multiplication, while 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 tokens, the embedding for the sequence is represented by a matrix of . 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 ) has a meaning in some high dimensional space called a “value” space (of size ). 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 should affect a token (where can be equal to ).

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 , and , which is calculated by 3 matrices , , and .

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

And each key query pair is resolved as:

Size of is noting the attention each token is paying to every other token.

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

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 ? #

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.

where is the row index and is the column index.

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

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 , while the output is of size . So either the rest of the layers need to take input of size or is set equal to and generally the latter is chosen but not for just this reason.

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 and at the final output location . Mathematically:

This does require

NNote on residual impact #

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

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

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 gets smaller and smaller for deeper layers.

Multi Head Attention (MHA)

Increasing efficiency

Attention mechanism represents each token as three 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 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 attention heads, each with internal representation of size . The output of each head is concatenated in the end to recover sized output vector of the attention layer.

In equations (where represents a specific attention head):

Output matrix

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

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

QWhy not keep the still in #

QWhy not merge with each head's as two transformations: ? #

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.

Summed up

So the overall equation for one transformer layer looks like:

Where and . After stacking layers, the output of layer () is normalized to produce the final transformer vector:

MHA matmaul #

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 , 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 product. In fact, it produces an additional term 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).

"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 where 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 . Now choosing a fixed rotation frequency and and representing the absolute position of those vectors, that term becomes .

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 real numbers into a high dimensional space of complex numbers.

For a 2D pair at position , the rotation by angle is:

Extending this to the full -dimensional embedding vector , the block-diagonal rotation matrix is defined as:

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

• If we choose a very fast moving frequency , then the dot product will start repeating very soon. Like for , , basically treating the token at 8 positions away as same as attending to self token.
• If we choose a very slow moving frequency , then the dot product will look almost the same. Like for , , basically failing to distinguish between nearby tokens.

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

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

Mathematically this looks like:

where the frequencies are .

Generalizing over the input sequence :

“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 . A reverse embedding (called “LM Head”) inverts that to a logits over the dictionary. So the first token is:

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 for all the tokens that we have already processed. In fact, their is not even required, just the to calculate how the new token is impacted by those past tokens. So, we just store that and 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 . 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.

Plugging in the specific parameters

These are the specifics of our Qwen model:

• Number of layers (): .
• RoPE base (): .
• Hidden size (): .
• Attention heads (): .
• KV heads: .
• Head dimension (): .
• FFN size (): .

Harness

A few more things to get Qwen to generate tokens:

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

1# /// script2# requires-python = ">=3.14"3# dependencies = [4#     "huggingface-hub>=1.27.0",5#     "jinja2>=3.1.6",6#     "numpy>=2.5.2",7#     "psutil>=7.2.2",8#     "transformers>=5.15.0",9# ]10# ///1112import argparse13from dataclasses import dataclass14import json15import os16import struct17import time18import numpy as np19import numpy.typing as npt20import psutil21from huggingface_hub import snapshot_download22from transformers import AutoTokenizer23from typing import Annotated2425EmbedxCtx = Annotated[npt.NDArray[np.float32], "shape=(d_embed, max_context)"]26EmbedxDict = Annotated[npt.NDArray[np.float32], "shape=(d_embed, dict_size)"]27ModelxEmbed = Annotated[npt.NDArray[np.float32], "shape=(d_model, d_embed)"]28EmbedxTokens = Annotated[npt.NDArray[np.float32], "shape=(d_embed, T)"]29ModelxTokens = Annotated[npt.NDArray[np.float32], "shape=(d_model, T)"]30ModelxModel = Annotated[npt.NDArray[np.float32], "shape=(d_model, d_model)"]31FFNxModel = Annotated[npt.NDArray[np.float32], "shape=(d_ff, d_model)"]32ModelxFFN = Annotated[npt.NDArray[np.float32], "shape=(d_model, d_ff)"]33KVxModel = Annotated[npt.NDArray[np.float32], "shape=(num_kv_heads * d_k, d_model)"]343536def get_ram_mb() -> float:37  process = psutil.Process(os.getpid())38  return process.memory_info().rss / (1024 * 1024)394041class WeightContext:42  def __init__(self, : "MemoryEfficientSafetensors"):43    self. = 44    self.loaded_tensors: list[np.ndarray] = []4546  def load(self, name: str) -> np.ndarray:47    arr = self..load_tensor(name)48    self.loaded_tensors.append(arr)49    return arr5051  def __enter__(self):52    return self5354  def __exit__(self, exc_type, exc_val, exc_tb):55    self.loaded_tensors.clear()565758class MemoryEfficientSafetensors:59  def __init__(self, repo_id: str = "Qwen/Qwen2.5-Coder-7B-Instruct"):60    print("Checking / down model weight shards from HuggingFace Hub...")61    self.model_dir = snapshot_download(62      repo_id=repo_id, allow_patterns=["*.safetensors", "*.json"]63    )6465    index_path = os.path.join(self.model_dir, "model.safetensors.index.json")66    with open(index_path, "r") as f:67      index = json.load(f)68    self.weight_map = index["weight_map"]6970    self.file_headers = {}71    shard_files = set(self.weight_map.values())72    for shard in shard_files:73      shard_path = os.path.join(self.model_dir, shard)74      with open(shard_path, "rb") as f:75        header_len = struct.unpack("<Q", f.read(8))[0]76        header_json = f.read(header_len).decode("utf-8")77        header = json.loads(header_json)78        self.file_headers[shard] = {"header_len": header_len, "header": header}7980  def load_tensor(self, name: str) -> np.ndarray:81    shard = self.weight_map[name]82    shard_path = os.path.join(self.model_dir, shard)83    info = self.file_headers[shard]84    meta = info["header"][name]8586    header_len = info["header_len"]87    start, end = meta["data_offsets"]88    shape = meta["shape"]89    dtype_str = meta["dtype"]9091    with open(shard_path, "rb") as f:92      f.seek(8 + header_len + start)93      raw_bytes = f.read(end - start)9495    if dtype_str == "BF16":96      u16 = np.frombuffer(raw_bytes, dtype=np.uint16)97      arr = (u16.astype(np.uint32) << 16).view(np.float32)98    elif dtype_str == "F32":99      arr = np.frombuffer(raw_bytes, dtype=np.float32)100    elif dtype_str == "F16":101      arr = np.frombuffer(raw_bytes, dtype=np.float16).astype(np.float32)102    else:103      raise ValueError(f"Unsupported tensor dtype: {dtype_str}")104105    return arr.reshape(shape)106107  def load_scope(self) -> WeightContext:108    return WeightContext(self)109110111@dataclass112class Qwen2_5_Coder_7B_Config:113  num_layers: int = 28114  hidden_size: int = 3584115  num_heads: int = 28116  num_kv_heads: int = 4117  head_dim: int = 128118  intermediate_size: int = 18944119  vocab_size: int = 152064120  rope_theta: float = 1000000.0121  rms_norm_eps: float = 1e-6122123124class Embedding:125  def __init__(self, w_embed: EmbedxDict):126    self.w_embed = w_embed127128  def embed(self, token_ids: list[int]) -> EmbedxTokens:129    return self.w_embed[token_ids].T130131132class QwenRotaryEmbedding:133  def __init__(self, head_dim: int, base: float = 1000000.0):134    self.head_dim = head_dim135    self.base = base136    inv_freq = 1.0 / (base ** (np.arange(0, head_dim, 2, dtype=np.float32) / head_dim))137    self.inv_freq = inv_freq138139  def apply(self, x: np.ndarray, positions: np.ndarray) -> np.ndarray:140    freqs = np.outer(self.inv_freq, positions)141    emb = np.concatenate([freqs, freqs], axis=0)142    cos = np.cos(emb)[None, :, :]143    sin = np.sin(emb)[None, :, :]144145    half = self.head_dim // 2146    x1 = x[:, :half, :]147    x2 = x[:, half:, :]148    rotate_half = np.concatenate([-x2, x1], axis=1)149    return (x * cos) + (rotate_half * sin)150151152def softmax(x: np.ndarray, axis: int = 0) -> np.ndarray:153  x_max = np.max(x, axis=axis, keepdims=True)154  exp_x = np.exp(x - x_max)155  return exp_x / np.sum(exp_x, axis=axis, keepdims=True)156157158class KVCache:159  def __init__(self):160    self.k_cache: dict[int, np.ndarray] = {}161    self.v_cache: dict[int, np.ndarray] = {}162163  def update(164    self, layer_id: int, new_k: np.ndarray, new_v: np.ndarray165  ) -> tuple[np.ndarray, np.ndarray]:166    if layer_id not in self.k_cache:167      self.k_cache[layer_id] = new_k168      self.v_cache[layer_id] = new_v169    else:170      self.k_cache[layer_id] = np.concatenate([self.k_cache[layer_id], new_k], axis=-1)171      self.v_cache[layer_id] = np.concatenate([self.v_cache[layer_id], new_v], axis=-1)172    return self.k_cache[layer_id], self.v_cache[layer_id]173174175class GroupedQueryAttention:176  def __init__(177    self,178    num_heads: int,179    num_kv_heads: int,180    d_model: int,181    w_k: KVxModel,182    k_b: np.ndarray,183    w_q: ModelxModel,184    q_b: np.ndarray,185    w_v: KVxModel,186    v_b: np.ndarray,187    w_o: ModelxModel,188    rope: QwenRotaryEmbedding,189  ):190    self.num_heads = num_heads191    self.num_kv_heads = num_kv_heads192    self.d_model = d_model193    self.d_k = d_model // num_heads194    self.queries_per_kv = num_heads // num_kv_heads195    self.w_k, self.k_b = w_k, k_b[:, None]196    self.w_q, self.q_b = w_q, q_b[:, None]197    self.w_v, self.v_b = w_v, v_b[:, None]198    self.w_o = w_o199    self.rope = rope200201  def forward(202    self, x: np.ndarray, positions: np.ndarray, kv_cache: KVCache, layer_id: int = 0203  ) -> np.ndarray:204    T = x.shape[1]205    q = self.w_q @ x + self.q_b206    k = self.w_k @ x + self.k_b207    v = self.w_v @ x + self.v_b208209    q = q.reshape(self.num_heads, self.d_k, T)210    k = k.reshape(self.num_kv_heads, self.d_k, T)211    v = v.reshape(self.num_kv_heads, self.d_k, T)212213    q = self.rope.apply(q, positions)214    k = self.rope.apply(k, positions)215216    k, v = kv_cache.update(layer_id, k, v)217    T_total = k.shape[-1]218219    k = np.repeat(k, self.queries_per_kv, axis=0)220    v = np.repeat(v, self.queries_per_kv, axis=0)221222    # Note: Modifying the matrix multiplication here to take advantage of vector maths for speed. We'll talk more about it in the inference post.223    scores = (q.transpose(0, 2, 1) @ k) / np.sqrt(self.d_k)224    causal_mask = np.triu(np.ones((T, T_total), dtype=bool), k=T_total - T + 1)225    scores[:, causal_mask] = -1e9226227    weights = softmax(scores, axis=-1)228    head_outputs = (v @ weights.transpose(0, 2, 1)).reshape(self.d_model, T)229230    return self.w_o @ head_outputs231232233class SwiGLUFFN:234  def __init__(self, w_gate: FFNxModel, w_up: FFNxModel, w_down: ModelxFFN):235    self.w_gate = w_gate236    self.w_up = w_up237    self.w_down = w_down238239  @staticmethod240  def swish(x: np.ndarray) -> np.ndarray:241    return x / (1.0 + np.exp(-x))242243  def forward(self, x: ModelxTokens) -> ModelxTokens:244    gate = self.w_gate @ x245    up = self.w_up @ x246    gated_act = up * self.swish(gate)247    output = self.w_down @ gated_act248    return output249250251class RMSNorm:252  def __init__(self, weight: npt.NDArray[np.float32], eps: float = 1e-6):253    self.weight = weight[:, None]254    self.eps = eps255256  def forward(self, x: np.ndarray) -> np.ndarray:257    variance = np.mean(x**2, axis=0, keepdims=True)258    return (x / np.sqrt(variance + self.eps)) * self.weight259260261class TransformerBlock:262  def __init__(263    self,264    attention: GroupedQueryAttention,265    ffn: SwiGLUFFN,266    norm1: RMSNorm,267    norm2: RMSNorm,268  ):269    self.attention = attention270    self.ffn = ffn271    self.norm1 = norm1272    self.norm2 = norm2273274  def forward(275    self,276    x: ModelxTokens,277    positions: np.ndarray,278    kv_cache: KVCache,279    layer_id: int = 0,280  ) -> ModelxTokens:281    norm_x1 = self.norm1.forward(x)282    attn_out = self.attention.forward(norm_x1, positions, kv_cache, layer_id)283    intermediate = x + attn_out284285    norm_x2 = self.norm2.forward(intermediate)286    ffn_out = self.ffn.forward(norm_x2)287    output = intermediate + ffn_out288    return output289290291class StackedTransformer:292  def __init__(293    self,294    : MemoryEfficientSafetensors,295    cfg: Qwen2_5_Coder_7B_Config,296    rope: QwenRotaryEmbedding,297  ):298    self. = 299    self.cfg = cfg300    self.rope = rope301302  def forward(303    self, x: ModelxTokens, positions: np.ndarray, kv_cache: KVCache304  ) -> ModelxTokens:305    for layer_idx in range(self.cfg.num_layers):306      prefix = f"model.layers.{layer_idx}."307      with self..load_scope() as ctx:308        in_norm_w = ctx.load(prefix + "input_layernorm.weight")309        norm1 = RMSNorm(in_norm_w, eps=self.cfg.rms_norm_eps)310311        q_w = ctx.load(prefix + "self_attn.q_proj.weight")312        q_b = ctx.load(prefix + "self_attn.q_proj.bias")313        k_w = ctx.load(prefix + "self_attn.k_proj.weight")314        k_b = ctx.load(prefix + "self_attn.k_proj.bias")315        v_w = ctx.load(prefix + "self_attn.v_proj.weight")316        v_b = ctx.load(prefix + "self_attn.v_proj.bias")317        o_w = ctx.load(prefix + "self_attn.o_proj.weight")318319        attn = GroupedQueryAttention(320          self.cfg.num_heads,321          self.cfg.num_kv_heads,322          self.cfg.hidden_size,323          k_w,324          k_b,325          q_w,326          q_b,327          v_w,328          v_b,329          o_w,330          self.rope,331        )332333        post_norm_w = ctx.load(prefix + "post_attention_layernorm.weight")334        norm2 = RMSNorm(post_norm_w, eps=self.cfg.rms_norm_eps)335336        gate_w = ctx.load(prefix + "mlp.gate_proj.weight")337        up_w = ctx.load(prefix + "mlp.up_proj.weight")338        down_w = ctx.load(prefix + "mlp.down_proj.weight")339        ffn = SwiGLUFFN(gate_w, up_w, down_w)340341        block = TransformerBlock(attn, ffn, norm1, norm2)342        x = block.forward(x, positions, kv_cache, layer_idx)343344    with self..load_scope() as ctx:345      final_norm_w = ctx.load("model.norm.weight")346      final_norm = RMSNorm(final_norm_w, eps=self.cfg.rms_norm_eps)347      output = final_norm.forward(x)[:, -1:]348349    return output350351352class LMHead:353  def __init__(self, w_head: npt.NDArray[np.float32]):354    self.w_head = w_head355356  def forward(357    self, last_hidden_state: npt.NDArray[np.float32]358  ) -> npt.NDArray[np.float32]:359    return (self.w_head @ last_hidden_state).squeeze(-1)360361362class Decoder:363  def __init__(364    self, transformers: StackedTransformer, : MemoryEfficientSafetensors365  ):366    self.transformers = transformers367    self. = 368369  def step(370    self, x: np.ndarray, positions: np.ndarray, kv_cache: KVCache371  ) -> tuple[int, np.ndarray, np.ndarray]:372    last_hidden_state = self.transformers.forward(x, positions, kv_cache)373    with self..load_scope() as ctx:374      lm_head_w = ctx.load("lm_head.weight")375      lm_head = LMHead(lm_head_w)376      logits = lm_head.forward(last_hidden_state)377    next_token_id = int(np.argmax(logits))378    return next_token_id, logits, last_hidden_state379380381def generate_tokens(382  max_new_tokens: int = 3, prompt: str = "Give me the quicksort algorithm in python"383):384  print("=== Qwen2.5-Coder-7B-Instruct pure NumPy Inference ===")385  print(f"Prompt: {prompt!r}")386  print(f"Max New Tokens: {max_new_tokens}")387  print(f"Initial Process memory consumption: {get_ram_mb():.1f} MB\n")388389  t_start = time.time()390391  tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen2.5-Coder-7B-Instruct")392  eos_token_ids = {393    t_id394    for t_id in [395      tokenizer.eos_token_id,396      tokenizer.convert_tokens_to_ids("<|im_end|>"),397      tokenizer.convert_tokens_to_ids("<|endoftext|>"),398    ]399    if t_id is not None400  }401  messages = [402    {403      "role": "system",404      "content": "You are a helpful assistant specializing in coding.",405    },406    {"role": "user", "content": prompt},407  ]408  text = tokenizer.apply_chat_template(409    messages, tokenize=False, add_generation_prompt=True410  )411  tokens = tokenizer.encode(text)412  seq_len = len(tokens)413  print(f"Formatted Chat Prompt:\n{text}")414  print(f"Tokenized Sequence Length: {seq_len} tokens.\n")415416   = MemoryEfficientSafetensors("Qwen/Qwen2.5-Coder-7B-Instruct")417  cfg = Qwen2_5_Coder_7B_Config()418419  rope = QwenRotaryEmbedding(head_dim=cfg.head_dim, base=cfg.rope_theta)420  transformers = StackedTransformer(, cfg, rope)421  decoder = Decoder(transformers, )422423  generated_tokens = []424  cache = KVCache()425426  for step in range(max_new_tokens):427    step_start = time.time()428    curr_seq_len = len(tokens)429430    if max_new_tokens > 1:431      print(432        f"\n--- Generating Token {step + 1}/{max_new_tokens} (seq_len={curr_seq_len}) ---"433      )434435    with .load_scope() as ctx:436      embed_weight = ctx.load("model.embed_tokens.weight")437      embedding = Embedding(embed_weight)438      if step == 0:439        x = embedding.embed(tokens)440        positions = np.arange(curr_seq_len, dtype=np.float32)441      else:442        x = embedding.embed([tokens[-1]])443        positions = np.array([curr_seq_len - 1], dtype=np.float32)444445    next_token_id, logits, output = decoder.step(x, positions, cache)446    next_token_str = tokenizer.decode([next_token_id])447448    top5_indices = np.argsort(logits)[-5:][::-1]449    top5_logits = logits[top5_indices]450    top5_tokens = [tokenizer.decode([idx]) for idx in top5_indices]451452    step_end = time.time()453454    tokens.append(next_token_id)455    generated_tokens.append(next_token_id)456457    print(f"Generated Token {step + 1} ID  : {next_token_id}")458    print(f"Generated Token {step + 1} Text: {next_token_str!r}")459    print(f"Top 5 Logits            : {top5_logits.tolist()}")460    print(f"Top 5 Tokens            : {top5_tokens}")461    print(f"Step Time               : {step_end - step_start:.2f} s")462    print(f"Process memory          : {get_ram_mb():.1f} MB")463464    if next_token_id in eos_token_ids:465      print(f"Reached EOS token ({next_token_str!r}, ID: {next_token_id}).")466      break467468  t_end = time.time()469470  full_generated_text = tokenizer.decode(generated_tokens)471  print("\n================ FINAL GENERATION RESULT ================")472  print(f"Total Tokens Generated  : {len(generated_tokens)}")473  print(f"Generated Tokens List   : {generated_tokens}")474  print(f"Generated Text          : {full_generated_text!r}")475  print(f"Total Computation Time  : {t_end - t_start:.2f} s")476  print(f"Peak Process memory: {get_ram_mb():.1f} MB")477  print("=========================================================\n")478479480if __name__ == "__main__":481  parser = argparse.ArgumentParser(482    description="Qwen2.5-Coder-7B-Instruct pure NumPy inference"483  )484  parser.add_argument(485    "num_tokens",486    type=int,487    nargs="?",488    default=None,489    help="Number of tokens to generate (positional)",490  )491  parser.add_argument(492    "--max-new-tokens",493    "-n",494    type=int,495    default=None,496    help="Number of tokens to generate",497  )498  parser.add_argument(499    "--prompt",500    "-p",501    type=str,502    default="Give me the quicksort algorithm in python",503    help="Input prompt",504  )505  args = parser.parse_args()506507  max_new_tokens = (508    args.max_new_tokens509    if args.max_new_tokens is not None510    else (args.num_tokens if args.num_tokens is not None else 3)511  )512  generate_tokens(max_new_tokens=max_new_tokens, prompt=args.prompt)5131---2[dependencies]3anyhow = "1.0"4hf-hub = "0.4"5rayon = "1.10"6serde = { version = "1.0", features = ["derive"] }7serde_json = "1.0"8sysinfo = "0.33"9tokenizers = "0.21"1011[profile.dev]12opt-level = 31314[profile.release]15opt-level = 316---1718#![allow(dead_code, unused_variables, deprecated, non_camel_case_types)]1920use anyhow::{Context, Result};21use hf_hub::api::sync::Api;22use rayon::prelude::*;23use serde::Deserialize;24use std::collections::{HashMap, HashSet};25use std::fs::File;26use std::io::{Read, Seek, SeekFrom};27use std::path::PathBuf;28use sysinfo::{Pid, System};29use tokenizers::Tokenizer;3031// --- MEMORY MONITORING ---3233pub fn get_ram_mb() -> f64 {34  let mut sys = System::new();35  let pid = Pid::from_u32(std::process::id());36  sys.refresh_processes_specifics(37    sysinfo::ProcessesToUpdate::Some(&[pid]),38    true,39    sysinfo::ProcessRefreshKind::nothing().with_memory(),40  );41  if let Some(proc_) = sys.process(pid) {42    proc_.memory() as f64 / (1024.0 * 1024.0)43  } else {44    0.045  }46}4748// --- MEMORY EFFICIENT SAFETENSORS  ---4950#[derive(Deserialize)]51struct IndexJson {52  weight_map: HashMap<String, String>,53}5455#[derive(Deserialize)]56struct TensorInfo {57  dtype: String,58  shape: Vec<usize>,59  data_offsets: (usize, usize),60}6162struct ShardHeader {63  header_len: u64,64  tensors: HashMap<String, serde_json::Value>,65}6667pub struct MemoryEfficientSafetensors {68  model_dir: PathBuf,69  weight_map: HashMap<String, String>,70  shard_headers: HashMap<String, ShardHeader>,71}7273impl MemoryEfficientSafetensors {74  pub fn new(repo_id: &str) -> Result<Self> {75    println!("Checking / down model weight shards from HuggingFace Hub...");76    let api = Api::new().context("Failed to initialize HuggingFace API client")?;77    let repo = api.model(repo_id.to_string());7879    let index_path = repo80      .get("model.safetensors.index.json")81      .context("Failed to get model.safetensors.index.json")?;8283    let model_dir = index_path84      .parent()85      .context("Failed to get parent dir of index file")?86      .to_path_buf();8788    let index_file = File::open(&index_path)?;89    let index: IndexJson = serde_json::from_reader(index_file)?;90    let weight_map = index.weight_map;9192    let unique_shards: HashSet<String> = weight_map.values().cloned().collect();93    let mut shard_headers = HashMap::new();9495    for shard in unique_shards {96      let shard_path = repo97        .get(&shard)98        .with_context(|| format!("Failed to fetch shard {}", shard))?;99      let mut file = File::open(&shard_path)?;100101      let mut header_len_bytes = [0u8; 8];102      file.read_exact(&mut header_len_bytes)?;103      let header_len = u64::from_le_bytes(header_len_bytes);104105      let mut header_bytes = vec![0u8; header_len as usize];106      file.read_exact(&mut header_bytes)?;107108      let tensors: HashMap<String, serde_json::Value> = serde_json::from_slice(&header_bytes)?;109      shard_headers.insert(110        shard,111        ShardHeader {112          header_len,113          tensors,114        },115      );116    }117118    Ok(Self {119      model_dir,120      weight_map,121      shard_headers,122    })123  }124125  pub fn load_tensor(&self, name: &str) -> Result<(Vec<f32>, Vec<usize>)> {126    let shard = self127      .weight_map128      .get(name)129      .with_context(|| format!("Tensor '{}' not found in weight map", name))?;130    let shard_path = self.model_dir.join(shard);131    let header_info = &self.shard_headers[shard];132    let raw_val = header_info.tensors.get(name).with_context(|| {133      format!(134        "Tensor '{}' not found in header for shard '{}'",135        name, shard136      )137    })?;138139    let tensor_meta: TensorInfo = serde_json::from_value(raw_val.clone())?;140141    let header_len = header_info.header_len;142    let (start, end) = tensor_meta.data_offsets;143    let byte_count = end - start;144145    let mut file = File::open(&shard_path)?;146    file.seek(SeekFrom::Start(8 + header_len + start as u64))?;147148    let mut raw_bytes = vec![0u8; byte_count];149    file.read_exact(&mut raw_bytes)?;150151    let shape = tensor_meta.shape.clone();152    let num_elements: usize = shape.iter().product();153154    let data = match tensor_meta.dtype.as_str() {155      "BF16" => {156        let mut arr = Vec::with_capacity(num_elements);157        for chunk in raw_bytes.chunks_exact(2) {158          let u16_val = u16::from_le_bytes([chunk[0], chunk[1]]);159          let f32_val = f32::from_bits((u16_val as u32) << 16);160          arr.push(f32_val);161        }162        arr163      }164      "F32" => {165        let mut arr = Vec::with_capacity(num_elements);166        for chunk in raw_bytes.chunks_exact(4) {167          let f32_val = f32::from_le_bytes([chunk[0], chunk[1], chunk[2], chunk[3]]);168          arr.push(f32_val);169        }170        arr171      }172      "F16" => {173        let mut arr = Vec::with_capacity(num_elements);174        for chunk in raw_bytes.chunks_exact(2) {175          let u16_val = u16::from_le_bytes([chunk[0], chunk[1]]);176          let f32_val = f16_to_f32(u16_val);177          arr.push(f32_val);178        }179        arr180      }181      other => anyhow::bail!("Unsupported tensor dtype: {}", other),182    };183184    Ok((data, shape))185  }186}187188fn f16_to_f32(h: u16) -> f32 {189  let sign = ((h >> 15) as u32) << 31;190  let exp = ((h >> 10) & 0x1f) as u32;191  let mant = (h & 0x03ff) as u32;192193  if exp == 0 {194    if mant == 0 {195      f32::from_bits(sign)196    } else {197      let mut m = mant << 1;198      let mut e = 0;199      while (m & 0x0400) == 0 {200        m <<= 1;201        e += 1;202      }203      let exp32 = (127 - 15 - e) << 23;204      let mant32 = (m & 0x03ff) << 13;205      f32::from_bits(sign | exp32 | mant32)206    }207  } else if exp == 31 {208    let exp32 = 0xff << 23;209    let mant32 = mant << 13;210    f32::from_bits(sign | exp32 | mant32)211  } else {212    let exp32 = (exp + (127 - 15)) << 23;213    let mant32 = mant << 13;214    f32::from_bits(sign | exp32 | mant32)215  }216}217218// --- OPTIMIZED PARALLEL MATRIX MULTIPLICATION WITH CHUNKING ---219220pub fn mat_mul(w: &[f32], m: usize, k: usize, x: &[f32], n: usize) -> Vec<f32> {221  let mut y = vec![0.0f32; m * n];222  let chunk_size = (m / (rayon::current_num_threads() * 4)).max(16);223224  y.par_chunks_mut(chunk_size * n)225    .enumerate()226    .for_each(|(chunk_idx, chunk)| {227      let start_row = chunk_idx * chunk_size;228      let rows_in_chunk = chunk.len() / n;229      for r in 0..rows_in_chunk {230        let j = start_row + r;231        let w_row = &w[j * k..(j + 1) * k];232        let y_row = &mut chunk[r * n..(r + 1) * n];233        for i in 0..n {234          let mut sum = 0.0f32;235          for idx in 0..k {236            sum += w_row[idx] * x[idx * n + i];237          }238          y_row[i] = sum;239        }240      }241    });242  y243}244245pub fn mat_mul_bias(w: &[f32], m: usize, k: usize, x: &[f32], n: usize, bias: &[f32]) -> Vec<f32> {246  let mut y = vec![0.0f32; m * n];247  let chunk_size = (m / (rayon::current_num_threads() * 4)).max(16);248249  y.par_chunks_mut(chunk_size * n)250    .enumerate()251    .for_each(|(chunk_idx, chunk)| {252      let start_row = chunk_idx * chunk_size;253      let rows_in_chunk = chunk.len() / n;254      for r in 0..rows_in_chunk {255        let j = start_row + r;256        let w_row = &w[j * k..(j + 1) * k];257        let b = bias[j];258        let y_row = &mut chunk[r * n..(r + 1) * n];259        for i in 0..n {260          let mut sum = 0.0f32;261          for idx in 0..k {262            sum += w_row[idx] * x[idx * n + i];263          }264          y_row[i] = sum + b;265        }266      }267    });268  y269}270271// --- KV CACHE FOR DECODING ---272273#[derive(Default)]274pub struct KVCache {275  // layer_idx -> (k_cache, v_cache)276  // k_cache: (num_kv_heads * head_dim, total_seq_len)277  pub cache: HashMap<usize, (Vec<f32>, Vec<f32>)>,278}279280impl KVCache {281  pub fn update(282    &mut self,283    layer_idx: usize,284    new_k: &[f32],285    new_v: &[f32],286    num_kv_heads: usize,287    head_dim: usize,288    seq_len: usize,289  ) -> (&[f32], &[f32], usize) {290    let kv_dim = num_kv_heads * head_dim;291    let entry = self292      .cache293      .entry(layer_idx)294      .or_insert_with(|| (Vec::new(), Vec::new()));295296    let prev_len = if entry.0.is_empty() {297      0298    } else {299      entry.0.len() / kv_dim300    };301    let new_len = prev_len + seq_len;302303    if prev_len == 0 {304      entry.0 = new_k.to_vec();305      entry.1 = new_v.to_vec();306    } else {307      let mut updated_k = vec![0.0f32; kv_dim * new_len];308      let mut updated_v = vec![0.0f32; kv_dim * new_len];309310      for d in 0..kv_dim {311        let src_old_k = &entry.0[d * prev_len..(d + 1) * prev_len];312        let src_new_k = &new_k[d * seq_len..(d + 1) * seq_len];313        let dst_k = &mut updated_k[d * new_len..(d + 1) * new_len];314        dst_k[..prev_len].copy_from_slice(src_old_k);315        dst_k[prev_len..].copy_from_slice(src_new_k);316317        let src_old_v = &entry.1[d * prev_len..(d + 1) * prev_len];318        let src_new_v = &new_v[d * seq_len..(d + 1) * seq_len];319        let dst_v = &mut updated_v[d * new_len..(d + 1) * new_len];320        dst_v[..prev_len].copy_from_slice(src_old_v);321        dst_v[prev_len..].copy_from_slice(src_new_v);322      }323      entry.0 = updated_k;324      entry.1 = updated_v;325    }326327    (&entry.0, &entry.1, new_len)328  }329}330331pub fn rms_norm(332  x: &[f32],333  hidden_size: usize,334  seq_len: usize,335  weight: &[f32],336  eps: f32,337) -> Vec<f32> {338  let mut out = vec![0.0f32; hidden_size * seq_len];339  let inv_rms: Vec<f32> = (0..seq_len)340    .into_par_iter()341    .map(|i| {342      let mut sum_sq = 0.0f32;343      for r in 0..hidden_size {344        let val = x[r * seq_len + i];345        sum_sq += val * val;346      }347      let mean_sq = sum_sq / (hidden_size as f32);348      1.0 / (mean_sq + eps).sqrt()349    })350    .collect();351352  out353    .par_chunks_mut(seq_len)354    .enumerate()355    .for_each(|(r, row)| {356      let w = weight[r];357      let x_row = &x[r * seq_len..(r + 1) * seq_len];358      for i in 0..seq_len {359        row[i] = x_row[i] * inv_rms[i] * w;360      }361    });362  out363}364365pub fn apply_rotary_pos_emb(366  q: &mut [f32],367  k: &mut [f32],368  seq_len: usize,369  start_pos: usize,370  num_heads: usize,371  num_kv_heads: usize,372  head_dim: usize,373  inv_freq: &[f32],374) {375  let half_dim = inv_freq.len();376377  for step_i in 0..seq_len {378    let pos = (start_pos + step_i) as f32;379    let mut cos = vec![0.0f32; half_dim];380    let mut sin = vec![0.0f32; half_dim];381    for i in 0..half_dim {382      let freq = pos * inv_freq[i];383      cos[i] = freq.cos();384      sin[i] = freq.sin();385    }386387    for h in 0..num_heads {388      let head_base = h * head_dim;389      for i in 0..half_dim {390        let idx1 = (head_base + i) * seq_len + step_i;391        let idx2 = (head_base + i + half_dim) * seq_len + step_i;392        let q1 = q[idx1];393        let q2 = q[idx2];394        q[idx1] = q1 * cos[i] - q2 * sin[i];395        q[idx2] = q2 * cos[i] + q1 * sin[i];396      }397    }398399    for h in 0..num_kv_heads {400      let head_base = h * head_dim;401      for i in 0..half_dim {402        let idx1 = (head_base + i) * seq_len + step_i;403        let idx2 = (head_base + i + half_dim) * seq_len + step_i;404        let k1 = k[idx1];405        let k2 = k[idx2];406        k[idx1] = k1 * cos[i] - k2 * sin[i];407        k[idx2] = k2 * cos[i] + k1 * sin[i];408      }409    }410  }411}412413pub fn softmax_inplace(x: &mut [f32]) {414  let max_val = x.iter().fold(f32::NEG_INFINITY, |a, &b| a.max(b));415  let mut sum = 0.0f32;416  for val in x.iter_mut() {417    *val = (*val - max_val).exp();418    sum += *val;419  }420  let inv_sum = 1.0 / sum;421  for val in x.iter_mut() {422    *val *= inv_sum;423  }424}425426pub fn compute_gqa_attention(427  q: &[f32],428  k: &[f32],429  v: &[f32],430  step_seq_len: usize,431  total_seq_len: usize,432  num_heads: usize,433  num_kv_heads: usize,434  head_dim: usize,435) -> Vec<f32> {436  let num_kv_groups = num_heads / num_kv_heads;437  let scale = 1.0 / (head_dim as f32).sqrt();438  let hidden_size = num_heads * head_dim;439440  let mut attn_out = vec![0.0f32; hidden_size * step_seq_len];441442  let head_outputs: Vec<(usize, Vec<f32>)> = (0..num_heads)443    .into_par_iter()444    .map(|h| {445      let kv_h = h / num_kv_groups;446      let mut h_out = vec![0.0f32; head_dim * step_seq_len];447448      let q_head_base = h * head_dim;449      let kv_head_base = kv_h * head_dim;450451      for i in 0..step_seq_len {452        let current_pos = total_seq_len - step_seq_len + i;453        let mut scores = vec![0.0f32; total_seq_len];454        for j in 0..total_seq_len {455          if j > current_pos {456            scores[j] = -1e9;457          } else {458            let mut dot = 0.0f32;459            for d in 0..head_dim {460              let q_val = q[(q_head_base + d) * step_seq_len + i];461              let k_val = k[(kv_head_base + d) * total_seq_len + j];462              dot += q_val * k_val;463            }464            scores[j] = dot * scale;465          }466        }467468        softmax_inplace(&mut scores);469470        for j in 0..total_seq_len {471          let weight = scores[j];472          if weight > 0.0 {473            for d in 0..head_dim {474              let v_val = v[(kv_head_base + d) * total_seq_len + j];475              h_out[d * step_seq_len + i] += weight * v_val;476            }477          }478        }479      }480      (h, h_out)481    })482    .collect();483484  for (h, h_out) in head_outputs {485    let q_head_base = h * head_dim;486    for d in 0..head_dim {487      let src_row = &h_out[d * step_seq_len..(d + 1) * step_seq_len];488      let dst_row =489        &mut attn_out[(q_head_base + d) * step_seq_len..(q_head_base + d + 1) * step_seq_len];490      dst_row.copy_from_slice(src_row);491    }492  }493494  attn_out495}496497pub fn swiglu(gate: &[f32], up: &[f32]) -> Vec<f32> {498  let mut out = vec![0.0f32; gate.len()];499  out500    .par_iter_mut()501    .zip(gate.par_iter())502    .zip(up.par_iter())503    .for_each(|((o, &g), &u)| {504      let silu = g / (1.0 + (-g).exp());505      *o = silu * u;506    });507  out508}509510// --- MODEL FORWARD PASS ---511512pub fn generate_tokens(max_new_tokens: usize, prompt: &str) {513  println!("=== Qwen2.5-Coder-7B-Instruct pure Rust Inference ===");514  println!("Prompt: {:?}", prompt);515  println!("Max New Tokens: {}", max_new_tokens);516  println!(517    "Initial Process memory consumption: {:.1} MB\n",518    get_ram_mb()519  );520521  let t_start = std::time::Instant::now();522523  let repo_id = "Qwen/Qwen2.5-Coder-7B-Instruct";524  let api = Api::new().unwrap();525  let repo = api.model(repo_id.to_string());526  let tokenizer_path = repo.get("tokenizer.json").unwrap();527  let tokenizer = Tokenizer::from_file(&tokenizer_path).unwrap();528529  let  = MemoryEfficientSafetensors::new(repo_id).unwrap();530531  let text = format!(532        "<|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",533        prompt534    );535  let encoding = tokenizer.encode(text.as_str(), true).unwrap();536  let mut tokens: Vec<u32> = encoding.get_ids().to_vec();537538  let initial_seq_len = tokens.len();539  println!("Formatted Chat Prompt:\n{}", text);540  println!("Tokenized Sequence Length: {} tokens.\n", initial_seq_len);541542  let num_layers = 28;543  let hidden_size = 3584;544  let num_heads = 28;545  let num_kv_heads = 4;546  let head_dim = 128;547  let rope_theta = 1000000.0f32;548  let eps = 1e-6f32;549550  // Precompute RoPE inverse frequencies once (mirrors QwenRotaryEmbedding).551  let half_dim = head_dim / 2;552  let mut inv_freq = vec![0.0f32; half_dim];553  for i in 0..half_dim {554    inv_freq[i] = 1.0 / rope_theta.powf((2 * i) as f32 / head_dim as f32);555  }556557  let eos_token_ids: HashSet<u32> = [558    tokenizer.token_to_id("<|im_end|>"),559    tokenizer.token_to_id("<|endoftext|>"),560  ]561  .iter()562  .filter_map(|&id| id)563  .collect();564565  let mut generated_tokens = Vec::new();566  let mut kv_cache = KVCache::default();567568  for step in 0..max_new_tokens {569    let step_start = std::time::Instant::now();570    let total_seq_len = tokens.len();571572    if max_new_tokens > 1 {573      println!(574        "\n--- Generating Token {}/{} (seq_len={}) ---",575        step + 1,576        max_new_tokens,577        total_seq_len578      );579    }580581    let (step_tokens, start_pos) = if step == 0 {582      (tokens.as_slice(), 0)583    } else {584      (&tokens[tokens.len() - 1..], total_seq_len - 1)585    };586    let step_seq_len = step_tokens.len();587588    // 1. Embedding Lookup589    let (embed_weight, _) = .load_tensor("model.embed_tokens.weight").unwrap();590    let mut x = vec![0.0f32; hidden_size * step_seq_len];591    for (i, &tok_id) in step_tokens.iter().enumerate() {592      let tok_idx = tok_id as usize;593      let src = &embed_weight[tok_idx * hidden_size..(tok_idx + 1) * hidden_size];594      for r in 0..hidden_size {595        x[r * step_seq_len + i] = src[r];596      }597    }598    drop(embed_weight);599600    // 2. Sequential Layer Forward Pass601    for layer_idx in 0..num_layers {602      let pfx = format!("model.layers.{}.", layer_idx);603604      // Pre-Attention RMSNorm605      let (in_norm_w, _) = 606        .load_tensor(&format!("{}input_layernorm.weight", pfx))607        .unwrap();608      let norm_x = rms_norm(&x, hidden_size, step_seq_len, &in_norm_w, eps);609      drop(in_norm_w);610611      // Projections612      let (q_w, q_shape) = 613        .load_tensor(&format!("{}self_attn.q_proj.weight", pfx))614        .unwrap();615      let (q_b, _) = 616        .load_tensor(&format!("{}self_attn.q_proj.bias", pfx))617        .unwrap();618      let mut q = mat_mul_bias(&q_w, q_shape[0], q_shape[1], &norm_x, step_seq_len, &q_b);619      drop(q_w);620      drop(q_b);621622      let (k_w, k_shape) = 623        .load_tensor(&format!("{}self_attn.k_proj.weight", pfx))624        .unwrap();625      let (k_b, _) = 626        .load_tensor(&format!("{}self_attn.k_proj.bias", pfx))627        .unwrap();628      let mut k = mat_mul_bias(&k_w, k_shape[0], k_shape[1], &norm_x, step_seq_len, &k_b);629      drop(k_w);630      drop(k_b);631632      let (v_w, v_shape) = 633        .load_tensor(&format!("{}self_attn.v_proj.weight", pfx))634        .unwrap();635      let (v_b, _) = 636        .load_tensor(&format!("{}self_attn.v_proj.bias", pfx))637        .unwrap();638      let v = mat_mul_bias(&v_w, v_shape[0], v_shape[1], &norm_x, step_seq_len, &v_b);639      drop(v_w);640      drop(v_b);641642      // RoPE643      apply_rotary_pos_emb(644        &mut q,645        &mut k,646        step_seq_len,647        start_pos,648        num_heads,649        num_kv_heads,650        head_dim,651        &inv_freq,652      );653654      // Update KV Cache655      let (full_k, full_v, cached_seq_len) =656        kv_cache.update(layer_idx, &k, &v, num_kv_heads, head_dim, step_seq_len);657658      // GQA Attention659      let attn_out = compute_gqa_attention(660        &q,661        full_k,662        full_v,663        step_seq_len,664        cached_seq_len,665        num_heads,666        num_kv_heads,667        head_dim,668      );669      drop(q);670      drop(k);671      drop(v);672673      // Output projection674      let (o_w, o_shape) = 675        .load_tensor(&format!("{}self_attn.o_proj.weight", pfx))676        .unwrap();677      let attn_proj = mat_mul(&o_w, o_shape[0], o_shape[1], &attn_out, step_seq_len);678      drop(o_w);679      drop(attn_out);680681      // Residual 1682      for i in 0..x.len() {683        x[i] += attn_proj[i];684      }685      drop(attn_proj);686687      // Pre-MLP RMSNorm & SwiGLU MLP688      let (post_norm_w, _) = 689        .load_tensor(&format!("{}post_attention_layernorm.weight", pfx))690        .unwrap();691      let norm_x2 = rms_norm(&x, hidden_size, step_seq_len, &post_norm_w, eps);692      drop(post_norm_w);693694      let (gate_w, gate_shape) = 695        .load_tensor(&format!("{}mlp.gate_proj.weight", pfx))696        .unwrap();697      let gate = mat_mul(698        &gate_w,699        gate_shape[0],700        gate_shape[1],701        &norm_x2,702        step_seq_len,703      );704      drop(gate_w);705706      let (up_w, up_shape) = 707        .load_tensor(&format!("{}mlp.up_proj.weight", pfx))708        .unwrap();709      let up = mat_mul(&up_w, up_shape[0], up_shape[1], &norm_x2, step_seq_len);710      drop(up_w);711      drop(norm_x2);712713      let mlp_act = swiglu(&gate, &up);714      drop(gate);715      drop(up);716717      let (down_w, down_shape) = 718        .load_tensor(&format!("{}mlp.down_proj.weight", pfx))719        .unwrap();720      let mlp_out = mat_mul(721        &down_w,722        down_shape[0],723        down_shape[1],724        &mlp_act,725        step_seq_len,726      );727      drop(down_w);728      drop(mlp_act);729730      // Residual 2731      for i in 0..x.len() {732        x[i] += mlp_out[i];733      }734      drop(mlp_out);735736      if max_new_tokens == 1 {737        if (layer_idx + 1) % 7 == 0 || layer_idx == 0 || layer_idx == num_layers - 1 {738          println!(739            "Layer {:2}/{num_layers} completed. Process RAM: {:.1} MB",740            layer_idx + 1,741            get_ram_mb()742          );743        }744      } else {745        if (layer_idx + 1) % 14 == 0 || layer_idx == num_layers - 1 {746          println!(747            "Token {}/{} - Layer {:2}/{num_layers} completed. Process RAM: {:.1} MB",748            step + 1,749            max_new_tokens,750            layer_idx + 1,751            get_ram_mb()752          );753        }754      }755    }756757    // 3. Final RMSNorm & LM Head Projection758    if max_new_tokens == 1 {759      println!("\nProcessing final norm and LM head...");760    }761    let (final_norm_w, _) = .load_tensor("model.norm.weight").unwrap();762    let x_norm = rms_norm(&x, hidden_size, step_seq_len, &final_norm_w, eps);763    drop(final_norm_w);764    drop(x);765766    let mut last_token_hidden = vec![0.0f32; hidden_size];767    for r in 0..hidden_size {768      last_token_hidden[r] = x_norm[r * step_seq_len + (step_seq_len - 1)];769    }770    let (lm_head_w, lm_shape) = .load_tensor("lm_head.weight").unwrap();771    let vocab_size = lm_shape[0];772    let logits = mat_mul(&lm_head_w, vocab_size, hidden_size, &last_token_hidden, 1);773    drop(lm_head_w);774775    let mut indexed_logits: Vec<(u32, f32)> = logits776      .iter()777      .copied()778      .enumerate()779      .map(|(idx, logit)| (idx as u32, logit))780      .collect();781    indexed_logits782      .sort_unstable_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));783784    let top5 = &indexed_logits[..5];785    let next_token_id = top5[0].0;786    let next_token_str = tokenizer787      .decode(&[next_token_id], false)788      .unwrap_or_default();789790    let top5_logits: Vec<f32> = top5.iter().map(|&(_, l)| l).collect();791    let top5_tokens: Vec<String> = top5792      .iter()793      .map(|&(id, _)| tokenizer.decode(&[id], false).unwrap_or_default())794      .collect();795796    let step_end = std::time::Instant::now();797    let step_duration = step_end.duration_since(step_start).as_secs_f64();798799    tokens.push(next_token_id);800    generated_tokens.push(next_token_id);801802    if max_new_tokens == 1 {803      let total_duration = step_end.duration_since(t_start).as_secs_f64();804      println!("\n================ FINAL GENERATION RESULT ================");805      println!("Generated Next Token ID  : {}", next_token_id);806      println!("Generated Next Token Text: {:?}", next_token_str);807      println!("Top 5 Logits            : {:?}", top5_logits);808      println!("Top 5 Tokens            : {:?}", top5_tokens);809      println!("Total Computation Time  : {:.2} s", total_duration);810      println!("Peak Process RAM Usage  : {:.1} MB", get_ram_mb());811      println!("=========================================================\n");812    } else {813      println!("Generated Token {} ID  : {}", step + 1, next_token_id);814      println!("Generated Token {} Text: {:?}", step + 1, next_token_str);815      println!("Top 5 Logits            : {:?}", top5_logits);816      println!("Top 5 Tokens            : {:?}", top5_tokens);817      println!("Step Time               : {:.2} s", step_duration);818    }819820    if eos_token_ids.contains(&next_token_id) {821      println!(822        "Reached EOS token ({:?}, ID: {}).",823        next_token_str, next_token_id824      );825      break;826    }827  }828829  let t_end = std::time::Instant::now();830  if max_new_tokens > 1 {831    let full_generated_text = tokenizer832      .decode(&generated_tokens, false)833      .unwrap_or_default();834    let total_duration = t_end.duration_since(t_start).as_secs_f64();835    println!("\n================ FINAL GENERATION RESULT ================");836    println!("Total Tokens Generated  : {}", generated_tokens.len());837    println!("Generated Tokens List   : {:?}", generated_tokens);838    println!("Generated Text          : {:?}", full_generated_text);839    println!("Total Computation Time  : {:.2} s", total_duration);840    println!("Peak Process RAM Usage  : {:.1} MB", get_ram_mb());841    println!("=========================================================\n");842  }843}844845fn main() {846  let args: Vec<String> = std::env::args().collect();847  let mut max_new_tokens: Option<usize> = None;848  let mut prompt: String = "Give me the quicksort algorithm in python".to_string();849850  let mut i = 1;851  while i < args.len() {852    let arg = &args[i];853    if arg == "-n" || arg == "--max-new-tokens" {854      if i + 1 < args.len() {855        max_new_tokens = args[i + 1].parse().ok();856        i += 1;857      }858    } else if arg == "-p" || arg == "--prompt" {859      if i + 1 < args.len() {860        prompt = args[i + 1].clone();861        i += 1;862      }863    } else if !arg.starts_with('-') && max_new_tokens.is_none() {864      max_new_tokens = arg.parse().ok();865    }866    i += 1;867  }868869  let max_new_tokens = max_new_tokens.unwrap_or(3);870  generate_tokens(max_new_tokens, &prompt);871}872

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)

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

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.

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 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 to project to some low-rank dimension which serves as a base for KV, generally called . This matrix is of size signifying the transformation from embedding size to size.
• A to project back to for key per head. This matrix is of size . - A to project back to for value per head. This matrix is of size .

For each token, its low-rank representation of size is stored in the cache. Whereas previously, two vectors per token per KV head of size (and generally like v/s ) 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 w.r.t. past tokens (say we have seen tokens and is the new low-rank kv cache of size ) for one head (omitting for clarity):

using matrix associativity

Notice the brilliance here, each does not have to be computed separately, rather the attention head constant matrix can be absorbed in to the newly computed query vector for each head, denoted by . 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 directly with .

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 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 only for the right logit and 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 is that the mean stays at and variance at 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 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 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 ( in Gemma 2 --- which BTW also used logit softcapping in attention layers as well rather than QK normalization with the value ). So, a hard cap is employed for softmax like .

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 is used, which has linear behavior within the limits and asymptotically approaches the limits. The formula used for LM head in Gemma 2:

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

source & further reading

layog.io — original article