LLM Speedrun: Architecture

A new three-part blog series by an unnamed author begins with 'LLM Speedrun: Architecture,' providing a hands-on guide to transformer architecture with runnable code examples in Python (using numpy) and Rust (using rayon). The series aims to explain LLM internals, starting with architecture and covering inference and training in subsequent posts.

AI Use Audience Python’s numpy https://numpy.org/ or with Rust’s rayon https://docs.rs/rayon/latest/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 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 https://en.wikipedia.org/wiki/Recursive neural network or their most successful derivative, LSTMs https://en.wikipedia.org/wiki/Long short-term memory . 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 https://arxiv.org/abs/1409.0473 . Folks at Google looked at it and said “maybe… maybe that’s all you need?” leading to the “Attention is all you need” https://arxiv.org/abs/1706.03762 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” rope . 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 residual . 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 why-pre-norm-over-post-norm . 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 decoding 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 input-setup , 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" https://huggingface.co/blog/designing-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” https://blog.eleuther.ai/rotary-embeddings/ . 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 , 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 - Loading 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, loader: "MemoryEfficientSafetensorsLoader" :43 self.loader = loader44 self.loaded tensors: list np.ndarray = 4546 def load self, name: str - np.ndarray:47 arr = self.loader.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 MemoryEfficientSafetensorsLoader:59 def init self, repo id: str = "Qwen/Qwen2.5-Coder-7B-Instruct" :60 print "Checking / downloading 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 loader: MemoryEfficientSafetensorsLoader,295 cfg: Qwen2 5 Coder 7B Config,296 rope: QwenRotaryEmbedding,297 :298 self.loader = loader299 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.loader.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.loader.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, loader: MemoryEfficientSafetensorsLoader365 :366 self.transformers = transformers367 self.loader = loader368369 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.loader.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 loader = MemoryEfficientSafetensorsLoader "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 loader, cfg, rope 421 decoder = Decoder transformers, loader 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 loader.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 LOADER ---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 MemoryEfficientSafetensorsLoader {68 model dir: PathBuf,69 weight map: HashMap<String, String ,70 shard headers: HashMap<String, ShardHeader ,71}7273impl MemoryEfficientSafetensorsLoader {74 pub fn new repo id: &str - Result<Self {75 println "Checking / downloading 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 loader = MemoryEfficientSafetensorsLoader::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, = loader.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, = loader606 .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 = loader613 .load tensor &format "{}self attn.q proj.weight", pfx 614 .unwrap ;615 let q b, = loader616 .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 = loader623 .load tensor &format "{}self attn.k proj.weight", pfx 624 .unwrap ;625 let k b, = loader626 .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 = loader633 .load tensor &format "{}self attn.v proj.weight", pfx 634 .unwrap ;635 let v b, = loader636 .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 = loader675 .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, = loader689 .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 = loader695 .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 = loader707 .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 = loader718 .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, = loader.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 = loader.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 transformer-architecture 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...