To put this in scale, the latest (again, at the time of this writing) models seem to be in the order of 500 billion parameters, with each of these parameters obviously being a number. This means that to run it we need to load 500 billion numbers, and not just that, we need to also perform computations on them (as was previously described in the post about the forward block)... that is a lot of numbers, and a lot of moving of numbers even without considering KV-caches and the like. This would mean that if we have a parameter count of $N$, and a parameter representation $b$ in (bytes per parameter), we could define $W$ as the total bytes used by the parameters.

So, to put numbers to it, if we stored each number as a single precision floating point number, we would be spending 4 bytes of memory per parameter, which would put us at around 2 [TB] of memory just to load the model into memory. Let's say we have that, for an (extremely coarse) approximation consider that we need to move all of those parameters for each token of input and generated, and assuming a bandwidth $B$ (in [B/s]), we could define our throughput as:

This is for $W$ total bytes of weights, so our throughput is our memory bandwidth $B$ divided by $W$. A high-end consumer PC has something like 90 [GB/s] for bandwidth, and some of the high-end dedicated AI GPU systems can go way higher with 8 [TB/s] per GPU. For those values we have (again, we are being super naive and coarse, but this is just to grasp the scale of the problem)

Neat, the gains from consumer to industrial grade hardware are definitely something, but 4 tokens per second is hardly an impressive number. It is a given that memory bandwidth will continue to increase, but that is generally a slow process (or at least slower than the rate at which models are being developed). So, if we want to get improvements soon, what's left is to fiddle with the other component, the number of bytes of weights per token.

Aside from the "DNN scale well with parameter count" claim, there's another that is quite relevant here: "DNN are quite tolerant to numeric differences in the parameters" (Dettmers et al.). This one is quite important because it allows us to run the model with way smaller precisions without losing quality in a noticeable way. We have two avenues for this: training the model in lower precision representations, and quantizing the model so inference can be performed with much decreased precision.

The first approach is self-explanatory, if a lower precision number is used to train, the model will have its weights encoded in that precision. This is a direct gain and we have this as a very common pattern that has been used for a while. For example, the Llama model we are playing with in this project, uses BF16 encoding, which is not-quite-half-precision (half precision has 1 sign, 5 exponent and 10 mantissa bits, where BF16 has 1 sign, 8 exponent and 7 mantissa).

Model quantization is also quite simple in its goal while its execution is a bit more esoteric (for an uninitiated like me :D). Essentially, the model checkpoint that was trained with some precision can be modified to operate at a lower precision. For example, a BF16 model can be modified to operate with 8-bit float or integer, or even 4-bit integer, which against the full-precision baseline in our example above would be an 8x improvement (and 4x against BF16). It should be noted that all our assumptions are gross oversimplifications akin to saying $\pi=10$, useful to convey an idea and its scale but far from accurate.

References​

LLM.int8(): 8-bit Matrix Multiplication for Transformers at Scale, Dettmers et al.

In terms of implementation, this is a somewhat simple thing to do, the requirements are (a) values need to be stored by head, and position; (b) when read we need to retrieve all the positions stored (this is the "we need $K$ and $V$ for all previous tokens"); and (c) we need values to be stored in a way that is read-efficient as we write once but read many times.

Let's start by describing the shape, the relevant tensors are: $Q$, $K$, and $V$. First, we define a few values for the dimensions:

• $h$: Heads
• $h_{kv}$: Key-Value heads.
• $d$: Head dimensions.
• $hidden$: Hidden dimension, essentially the embedding dimensions.

Then how we compute each of them (and their dimensions):

Now, these vectors are the tensors for all heads/kv-heads, so each of them is of this form:

Given this, going forward we'll use $q_i$, $k_i$, and $v_i$ to denote the specific tensor for a head/kv-head.

Next, for how they are used, for each token at position $pos$ we do (note that for simplicity we will use $h$ to denote the appropriate kv-head too, see the previous post to see how we get that ;)):

Which is equivalent to:

From this it is clear that, for each token we compute its $k_h$ and we also need the values from every token up to there, so we just store it:

And the retrieve method should give us a shape that we can use in the vectorized version of the $scores_h$ computation, so:

Cool, now to the $V$ part: the usage of $v_h$.

Which is a bit of a mouthful, but can be rewritten as:

So, this is a little bit funnier. We have the option of writing the double loop or transpose and use our matvec kernel. Initially I did implement the double loop, but later updated to the other alternative (which is what we'll use here). It's clear that we need all of the old $v_h$, so we do the same as we did for $k_h$ but we store transposed:

Interface has the same shape, although less efficient to write for a row-aligned matrix (which is what we use), not terribly worrying as we write very few times compared to how many times we read.

And that is actually all there is to it, exactly what you would expect of something that has cache in the name, with the added interesting bit that shape matters.

Next up, we'll talk about floating-point numbers.

We'll start just after the application of $RoPE$. We have embedded position into our query ($Q$) and keys ($K$), the important clarification to do here is that we treated these tensors as vectors, but in reality they are a collection of vectors that are packed as a single vector, so for $Q$:

Where $h$ is the number of heads and $d$ is the head dimension.

$K$ and $V$ (the value vector) have a different layout from $Q$ (both are identical), but here each row will have the same dimension $d$ and there will be a different number of rows, the called "key-value heads" or $KVH$ for friends. The relationship between number of heads and the key-value heads is a model design decision and can define the following taxonomy:

• Multi-head Attention (MHA), same number of heads and key-value heads; the original transformer, good quality but more computationally expensive (this will be apparent in the next few paragraphs)
• Multi-query Attention (MQA), single key-value head is shared across all heads; this is cheaper to compute, but apparently has instability and quality issues.
• Grouped-query Attention (GQA), multiple key-value heads but they are each shared by a subset of the attention heads (e.g Llama 3 8B has 32 heads and 8 key-value heads)

Now that we have that out of the way, we'll enter into what is done for each attention head. First though, we need to determine which key-value head interacts with the current head. Let's say we are processing attention head $i$ ($0 \leq i < h$), we then will determine the key-value head $j$:

We reached the part where in the previous post we said something to the effect of "we then apply our query to the keys of all previous tokens" and we also magically introduced a $i$ (master foreshadowing). Untangling all of the magic, we'll now say, we get all previously computed $K$ vectors corresponding to head $i$ from the cache (ok, not all of the magic, some mirrors left) and for computation convenience we'll pack them side by side to build a matrix (let's say current token is number $c$):

We then proceed to compute the scores that correspond to this head:

Then these scores need to be scaled. The explanation for the scaling and its value is that this product adds the components of all head dimensions. Here the assumption made is that each of these components is an independent variable, and if its mean is $0$ and variance is $1$ it follows that the variance of the sum is $d$ ($d$ dimensions, $d$ components in the sum), so scaling by the standard deviation would normalize (standard deviation is $\sqrt{d}$), so:

Now that scores are standardized, we use $\text{softmax}$ to turn them into probabilities:

And we use these to compute the attention for this head by applying it to the values of all tokens so far (we obtain $V_j$ in a way identical to $K_j$)

Once this is done, we just have to project this into something we can add to the residual stream, and for this we use the $W_O$ weights. This is somewhat interesting too, as these weights are packed as a single matrix but in reality are $h$ matrices.

This is why we concatenate all of the heads attention results into a single vector of size $d \cdot h$, so we can do:

At this point, we can zoom out and continue moving forward in the block as described in the other post (add the projection to the residual stream, and move to the FFN part).

References​

• GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints, Ainslie et al.
• Attention Is All You Need, Vaswani et al.

The format​

First let's talk about the format, safetensors is a file format designed for storing tensors that claims safety and speed are its big advantage over its competitors. Apparently, some of the existing formats and their associated machinery execute code on load (pickle), and/or create copies of the tensors in sort of a two step load; compared to that, safetensors indeed can claim an advantage and the format is just that, a format: passive data waiting to be read, and data that is packed in a way that we can directly use once loaded into memory (the community seems to use zero-copy when referring to this property).

Now, concretely for the format, the layout is like this:

The header size at the beginning is just a LE uint64, and as the name says, it gives us the size of the next block, that is a json format header. The schema of this header is (using TypeSpec for simplicity):

union DataType {
  "BF16",
  "F64",
  "F32",
  "F16",
  "I64",
  "I32",
  "I16",
  "I8",
  "U8",
  "BOOL",
}

model TensorDescription {
  /**
   * Data type of the tensor
   */
  dtype: DataType;

  /**
   * Shape conveys the extents of the tensor.
   * A vector would have a single dimension  [D]
   * A matrix would have two dimensions [R, C]
   * An n-dimensional matrix would be [D_1, ..., D_N]
   */
  shape: uint64[];

  /**
   * The offsets where the tensor can be found,
   * of the form [start, end).
   * (relative to the end of the header).
   */
  @maxItems(2)
  @minItems(2)
  data_offsets: uint64[];
}

/**
 * Structure that contains the declaration of all tensors. Every
 * key is a tensor name, with its value being the description of
 * that tensor. The only exception is the `__metadata__` field that
 * is a string to string map with extra information.
 */
model SafetensorsHeader {
  __metadata__?: Record<string>;
  ...Record<TensorDescription>;
}

An interesting note about the header block is that it can be padded with trailing spaces which is relevant for alignment. Similarly, tensors are not necessarily densely packaged to account for alignment.

The parser​

To read the file in an easy way, we use mmap to efficiently map the file to memory. To initialize the mapping, we need the size of the file (that we can easily get with std::filesystem::file_size), and a file descriptor for the file (easy again with fopen and fileno, it should be noted that we can close the file stream with fclose as soon as we initialize the map as it does not require it to be open). Additionally, we pass the PROT_READ for the memory protection parameter (you may have guessed that means "pages can be read"), and the MAP_PRIVATE flag, essentially saying that the map is process local with copy-on-write (not relevant for us because read-only) It should be noted that since RAII is our friend, we use it for the closable stuff to be safe, homebrew scope_guard for the file stream, and std::shared_ptr with a custom deleter for the mmap (munmap called by the destructor lambda).

Now that we have the entire file as a big chunk of (virtual) memory, we can just start reading it. First 8-byte word is the size of the header. With it we can initialize a std::string_view that encompasses the whole header, and we use that to initialize our header struct.

For json deserialization, we use glaze. glaze is a really cool library that is amazingly fast at deserializing and allows us to have a type safe representation of the struct (no json["__metadata__"] like other libraries). It also allows to easily customize the parser and do stuff like "we only know this property, put the rest in this map" (which is really useful when that's exactly the structure of the header)

With the header at hand, we know all the tensors by name, and what their shape is, so we can now just get them. For this we build a thin wrapper class that has a view of the memory scoped to the tensor position and size within the memory map. This class also has an alias std::shared_ptr to guarantee lifetime of the memory map is tied to its lifetime, and, of course, information about its shape.

To ensure everything works as intended, I asked the AI to help me generate a python script that would generate a sample safetensors file and a header file with the expected values. To accompany these, I also asked for a test file to exercise my code (with this being the part I wasn't as interested in writing, having it help with this ended up being really useful)

The code for this part of the project (with some extra stuff where I was reorganizing some of the code) can be found here.

To understand how these operations are made and how information is accumulated across a pass to all the model layers, we need to understand what the literature calls the "residual stream". The residual stream starts as the original embedding (the vector representation of the token), and as we move forward in the block we add all of the results (this can be imagined as the block imbuing the original vector with meaning derived from context and the model's training), and so, it is "residual" because each layer adds its contribution to the running result (always of the form $x_{i+1} = x_i + f(x_i)$).

Now, in terms of what actually happens, the first (arguably trivial) operation is to transform the received token (a number) into its equivalent embedding. This is done by doing a simple table lookup on the vocabulary table that ships with the model; the token is just the index of the row that contains the vector.

With the embedding in hand, we first normalize it (as it was mentioned in the previous post this is to mitigate the drift that the compound operations can cause in our residual stream). Note that we normalize into a different vector, so we keep our residual vector residual (in notation we'll include the weights for normalization but will elide the $eps$ parameter).

Once normalized, we use it to generate the values we require for the attention. Attention is the mechanism by which the model imbues context into the embedding of a specific token. For this, we need to determine how tokens affect each other, and also in which way. The model provides three sets of weights (per layer, everything here is per layer) to obtain exactly this information:

The query ($Q$) and key ($K$) vectors are used to determine how a token attends another, and the value ($V$) is the content that gets modulated by that attention. Something that is interesting about these values, is that they are split across several parallel heads (this is something that we'll punt to another post to keep this one reasonable)

Next part is imbuing position information into $Q$ and $K$. We do that with RoPE:

Once this is done, we need to apply our query $Q$ to the key $K$ of all previous tokens to compute how they influence each other, with this yielding the scores.

These scores get then scaled by the inverse of the square root of the dimension of each attention head, and softmax gets applied to it to get probabilities from the logits.

Now that we have built the scores as probabilities, we use them to calculate the contribution from $V$ that should actually be included in the attention, so for each attention head we have:

And with all heads having produced their attention vectors, we can concatenate them into a full width vector that we can then project using the $W_O$ weights.

And the final step of the attention block is to add the projection to the residual stream ($x_{i+1} = x_i + f(x_i)$, ta da!):

The next part we then run within the layer is what literature generally calls FFN (Feed-forward Network) or sometimes MLP (Multi-layer Perceptron). This part involves a linear part, a non-linear part that is generally called activation, and a final linear part that produces the results that are then added to the residual stream. This is the second and last part of the processing of each layer and differs from attention in that semantic enrichment comes not from the neighboring tokens, but rather the model enriches each token independently (this is of course a useful mental model, but it's hard to track what's going on exactly here).

As in the attention layer, the first thing we do is to normalize the input.

Once normalized, we need to generate a couple of projections in what is the first linear phase mentioned above:

Where $U$ (up) is the vector that contains the meaning that is "trying" to make it to the residual stream, and $G$ (gate) is the vector that determines how the activation of that happens (note that both $G$ and $U$ have more dimensions than $H$ so $W_{\text{up}}$ "lifts" $H$). First, we build the SiLU gate:

And then it just gets applied:

And finally, we apply the $W_{\text{down}}$ weights to project the output (bring down into the dimensions of $H$):

And to close (double ta da!), the result gets added to the residual stream:

This concludes the layer part of the block, the above transformations get applied for each layer of the model (each layer has its own set of weights) with the results always accumulated in the residual stream.

Once the layer processing is complete, we are ready to compute the next token from the residual stream. To do this, we first normalize:

Then we apply the unembedding (lm_head) to produce the logits.

And finally, in our greedy implementation we get the token by using argmax:

If we are processing input, we discard this token and process the next input token. This does not mean the pass was a waste as it updates the KV-cache (topic that we punt to a future post). If we are in the generation phase (past the input), this would be our next output token. GG

References:​

• A Mathematical Framework for Transformer Circuits
• 3Blue1Brown: Transformers, the tech behind LLMs
• 3Blue1Brown: Attention in transformers, step-by-step
• 3Blue1Brown: How might LLMs store facts

This being a learning adventure, the progression I arrived at was: first, I'll write a naive implementation that lets me learn the process, and we'll refine from there.

The components to get the naive implementation to work are:

• The operations required (kernels) to run generation.
• The KV cache (critical part that keeps past calculations so each token considers them without recomputing)
• The structure that loads and holds the model weights and config.
• The actual code that runs generation.

Tokenization in both directions is for now out of scope and I'll use a wrapping layer that handles token encoding/decoding. On top of that, we have a somewhat orthogonal staging for implementation that relates to runtime optimization, and note that we will start with the slowest most naive implementation with the explicit intention of experiencing the impact of each layer. These, at broad strokes are: scalar implementation, SIMD, concurrency, and GPU.

Finally, the last axis in this implementation is related to how it works (and how useful it is). Broadly speaking we have the greedy implementation (argmax) vs the non-deterministic (with top-p and temperature). This is one of the most interesting parts from a learning perspective as it is the one that I understand the least, so for now we just parrot the names.

With that, let's move to the first stage, the kernels. From the writeup AI generated for me, it seems that we will need to implement a few operations (these are presented as refdoc-like signatures sometimes with pseudocode). These are:

• embed(table: [r, c], token: int)- > [c]

  Essentially, a table lookup where a token_id (index) is received and a vector is returned

• rmsnorm(x: [r], w: [r], eps: float) -> [r]

  Root mean square normalization. Apparently, since each layer computes a number of multiplications, the impact of these start compounding and the values can drift. To mitigate this LayerNormalization is introduced, and RMS Normalization is a variant on the traditional version of it that is cheaper computationally.

• matvec(w: [r, c], x: [c]) -> [r]

  Plain matrix vector multiplication, easy enough.

• rope(t: [r, c], pos: int, theta_base: float) -> [r, c]

  Huh? This one needed learning, so watched a couple of videos on it. At its most basic, RoPE encodes position information in a token by rotating it. So here we rotate the vectors in t to encode the position pos.

  The reading yielded that this uses the fact that rotating vectors by an angle proportional to the position makes the dot product between two vectors depend only on their relative distance within the attention.

• softmax(x: [r]) -> [r]

  Again, something that needed a bit of reading. Apparently this function is used to turn a vector of n scores into a probability distribution. It's used to transform the attention scores into weights to model the relationship between tokens, and also to compute which token is the most likely next token (this happens at the end of the forward pass for each token)

• add(a: [r], b: [r]) -> [r]

  Just vector addition.

• silu_mul(gate: [r], up: [r]) -> [r]

  The activation function, this uses Swish on the gate and then it multiplies by up ultimately having the gate modulate how much of up goes past it (the SwiGLU gate)

• argmax(x: [r]) -> int

  Returns the index of the biggest element, this we use in our greedy implementation to find the next token.

This is what my research yielded for each of the functions I needed to implement and, while this has been really illustrative, I feel that I lack the conceptual understanding on how these play together. I'll shelve this concern for now and probably take a little side quest to find out more for the next writeup.

In terms of implementation, first thing to note is that none of these allocate, the memory for the output will always be provided by the caller. Second, some of the signatures don't show it, but operations are made in place where it makes sense.

Now, more deeply into the details, I decided to use std::mdspan for matrices and std::span for vectors (this is one of those decisions that a voice in my head tells me I'm going to regret it, but I ignored it to get something working first)

Last bit of the implementation here, AI provided a test file, and a python script to generate the golden values we would be testing against.

This first implementation can be found here (a number of updates were done after that, including acting on the advice to rethink the std::span decision)

To put a spin on this, I'm starting to experiment with using AI as a non-coding advisor, with the flow being: I select a topic, use AI to scope and frame what needs to be done, and start developing. Note that by non-coding, I mean not coding towards the problem, but still using it to help with mechanical parts like scaffolding. And the advisor part comes from reviewing, and pointing out conceptual gaps I need to close to understand the problem.

As a first attempt, I decided to implement an engine that can run LLMs in C++ from scratch. This seemed like a good project because a) it has been a few years since I started using these tools, and I do feel I don't understand how they work underneath, and b) C++ is the language I feel most comfortable with, and it's also a language I don't get to use in my day-to-day work (so this is my fun plug).

Without going into prompts and details, the first thing I did was to iterate over what I wanted to do and, importantly, how. The "how" here carries the explicit ask for "don't give me code", "review, but don't update", and so on. It also carries the progression I want from the project: small-ish tasks with observable (testable even) results, that build into an end-to-end that works. Once we reach "works", the real fun begins as we begin working towards "works well" and "works fast".