LPU Lite

01

What is an LPU?

LPU stands for Language Processing Unit. It is a specialized computer chip developed by Groq¹ specifically for running large language models such as ChatGPT and Claude. Groq has kept the secret of their LPU closed-source and only left small hints in their research paper.^{2, 3} The LPU differs from a traditional chip in many ways, but the overarching idea is “deterministic execution.” Deterministic execution refers to the idea that the chip does exactly what it is told and nothing else, with predictable timing and repeatable behaviour. If an operation is executed multiple times, it should result in the same number of clock cycles⁴ for each execution. Hardware is dumb. A chip does not think. It’s just a collection of microscopic physical switches (transistors) arranged to route signals. The philosophy behind the design of the LPU strips complex control logic off the chip and moves all the intelligence to a piece of software called the compiler.

02

Why did we build LPULite & what was our goal?

We had no guide or course that teaches chip design at our university. We had taken a digital logic course, but were disappointed with the fact that the most complex project we did was building a full adder in Quartus using logic blocks, not even in RTL!⁵ Therefore, we decided to challenge ourselves to dive deep into machine learning (ML) hardware and learn as much as we could on our own. We wanted to prove that basic math (like y = mx + b) and basic logic circuits are enough to help anyone understand how modern AI hardware works.

Our goal was to design our own version of the LPU from scratch and run a simple Transformer-style model on it, proving that with minimal Machine Learning and computer design knowledge, it’s totally possible. We were also driven by a simple question: What makes the LPU architecture so compelling that even Nvidia licensed it?

Keep in mind, this article is not intended to serve as a tutorial for “how to build an LPU from scratch,” and our architecture is not a 1:1 LPU. It serves as an educational resource for how someone with minimal hardware experience can approach this field, and our journey in building what we think an LPU would look like.

03

Our process

Realizing that in our day-to-day lives, we had been relying heavily on AI, we wanted to build something where we maintained ownership of our own thinking and design. We acknowledge that AI is used widely in today’s age to accelerate processes, but we wanted to use it as a tool to support our work rather than replace our learning or understanding.

Each time we tackled a problem, our initial step was drawing it out. From the most minor details to big ideas, drawing always got us thinking and helped us visualize the steps to take to solve them. From there, we verified that our ideas fit the philosophy of an LPU and moved on to minimal coding exercises (i.e getting AI to build us a skeleton of a Verilog⁶ module with inputs and outputs described, and we just have to fill in the middle!) to build a version of the task. This helped us build a great understanding of how everything worked right from the get-go. A significant portion of our AI usage was to understand what was actually done by reading the code and analyzing whether the design was implemented to our perfection.

04

Background on AI accelerators

Central Processing Units (CPUs) are chips you will see in every laptop; however, they were designed for general-purpose use and executing functions that make a computer work, but never for intense arithmetic workloads. They also work sequentially, completing tasks one by one. Moreover, a large portion of a CPU's silicon is dedicated to memory cache hierarchies and dynamically optimizing the timing of instructions. This becomes a bottleneck for ML arithmetic workloads.

Graphics Processing Units were created for rendering graphics (hence the word “graphics” in their name). Rendering graphics takes a lot of intense arithmetic computations; therefore, they were also adopted for machine learning! They introduced the concept of parallel arithmetic (computing thousands of math operations at once).

Dedicated accelerators like TPUs, NPUs and LPUs are designed specifically for ML arithmetic (not repurposed or general-purpose). Because they are so specified, they dedicate all their silicon strictly for ML to dramatically increase processing speed and improve power efficiency.

05

Background on hardware

Prior to starting this project, we thought a project related to building hardware meant we had to deal with physical objects, such as an Arduino board, printed circuit boards, soldering, etc. It turns out that for digital logic, you have to describe the hardware using code similarly to how you’d design a printed circuit board with computer-aided design! For example, let’s say you are coding a module to add two 8-bit numbers together to get a 9-bit number; here is how you would imagine that in hardware.

In this example, we have two input wires that represent two 8-bit inputs, and an output wire representing a 9-bit output. The “module” in between those wires represents our actual logic of how those two numbers are added together.

This logic in SystemVerilog can be represented as:

module adder (
    input  [7:0] a,
    input  [7:0] b,
    output [8:0] sum
);
assign sum = a + b;

endmodule

Now, to take this from software to hardware, you have to use something known as an FPGA (Field Programmable Gate Array). An FPGA can be used to test that your digital logic actually works in hardware; this is known as synthesis. For this project, we will be using the DE1-SoC board to test our chip, and Quartus to synthesize our design.

06

The Model We Targeted: MicroGPT

MicroGPT is a scaled-down Transformer that is trained on simple text, like names, to predict the next character in the sequence. Our focus was not on scale, but on validating the core concept. If we can prove why a module can run a certain arithmetic operation efficiently, we can scale it up given the tools we have today. However, the most important part remains proving why something works. By targeting MicroGPT, we got rid of any unnecessary parameters and isolated the most important idea of the entire project.

Below is an interactive demo to showcase how MicroGPT predicts the next character in a name. Click on 1 of the 4 letters to “generate” a name.

07

What is a Transformer?

Imagine reading a sentence where words contextually highlight each other. In the sentence "The bank of the river," the word "bank" connects to "river" rather than "money." That is the idea behind a Transformer: to see what words “attend” to others. To actually generate text, transformers “predict” what the next word is going to be in something called Next-token generation. Next-token generation is the process of predicting the very next word based on previous context, repeating this loop over and over. For example, when you type on your phone and the keyboard suggests the next word, that is the transformer at work predicting what you are going to say next. Therefore, a transformer's job is to predict what is going to be said next. Transformers are the backbone of all modern LLMs like GPT and Claude that we all use.

Before we dive into the actual hardware, we need to understand the inner workings of how language models figure out what words matter to each other in a sentence.

This is the formula that teaches AI which words to “pay attention to.”

The Q, K, V variables in this formula represent the Query, Key, and Value vectors.

Now let's take the input “Thinking Machines.” If we look at the words individually, who is thinking? What is being thought about? Who is the machine? What does the machine do? These questions can be answered through the QKV vectors!

In terms of Query, the word “Thinking” may encode the questions: what is the subject? What is the object? What are the nearby nouns? The Query vector includes these questions so the model can learn which information is relevant.

For Key, the word “Machines” may encode “I am a noun,” “I can be an agent,” “I often act.” So, K is like labels for each word. Therefore, Query searches through Keys to find answers.

As for Value, this vector may encode the information you get when you choose this word; V decides how much context you take per word.

These matrices are generated by multiplying the input (X) with Query weights, Key weights, and Value weights.

These weights are generated during training, which is a whole different part of machine learning; however, for now, we will only look at what happens during inference (using the model).

Each input (“dog”, “bites”, “man”, all of these are independent inputs) has its own Q, K, V vectors.

Once we have the Q, K, V vectors, we can then multiply \(Q @ K^{\mathsf T}\). Why a multiplication? If we look at the dot product and why it's done, the purpose is to understand how close two different vectors are. The purpose of Q is the question the word is asking.

If we take the Q vector for the machines, and matmul it with the K vector of machines, it will give a value very close to 1.00 indicating the two vectors are very similar, but if we matmul the Q for machines, and K for printer we may get a non zero vector, also not close to one, this would mean “the vectors have some relationship to one another, but are not the same.” The matmul of Q for “dog” and K for “man” might be way closer to the previous one, as it might denote “the vectors are both machines used for some sort of processes.”

The purpose of the matmul in this case is to give the vectors some sort of relationship to one another! We then follow up with scalar/vector operations to complete this formula.⁷

The vast majority of arithmetic work in a standard transformer consists of matrix multiplications (matmuls), making matmul acceleration critical for modern AI accelerators. Inside the LPU, the Matrix Execution Module (MXM) is responsible for computing matmuls efficiently.

08

MXM

The first step to designing our MXM was to understand the anatomy of a matmul. At its core, the matmul consists of many independent dot products between rows and columns of two different matrices. Each dot product requires a sequence of multiplications, followed by additions, which can be implemented through a multiply-accumulate (MAC) unit.

The MAC unit then becomes the fundamental computational block of our MXM.

Now, to zoom out and understand how these dot products are arranged in a matmul, let's look at an example for a 2x2 @ 2x2 matmul, and write out the resultant dot products.

We notice that the dot products are also in an array, and if each MAC unit is responsible for one dot product, we can arrange these MAC blocks into something known as a MAC array!

To understand how the MAC array works, let's look at the first outer product of the matmul:

When we place initial outer products inside the mac array, we get this arrangement:

Over here, we can observe from the topology of the matmul that the rows of the MAC array share inputs (highlighted in green), and the columns share the weights (highlighted in red). As a result, when we wire the array, we can have a row of the MAC share the input wire, and a column share a row wire.

A MAC array also allows us to introduce “parallelism,” meaning each MAC unit can compute the dot product and accumulate independently!

When we introduce the wires into the MAC array, it can be visualized as:

In the diagram above, the green wires carry the inputs to the right, and the red wires carry the weights down a column.

This is how we first visualized the MXM! We then realized that for this arrangement, inputs would need to pass through a MAC unit before reaching downstream MACs. Meaning, if we want i00 to go to MAC 0,1, it would have to go through MAC 0,0 before reaching its destination. It would mean that for MAC 0,1 to compute its respective dot product, it would be dependent on MAC 0,0’s timing (if MAC 0,0 is not complete, it cannot go to MAC 0,1).

We decided that the inputs and weights should arrive at their respective MACs at the same time to maintain our idea of independence. That would cause the MAC array to look something closer to this visualization:

This tweak in the design allowed us to bring our idea of inputs and weights arriving at their respective MACs without delays. This style of wiring also introduces a potential flaw in the design. The wiring gets increasingly more complicated as the design scales up. This is only a 2x2 MXM, whereas the real LPU is a 320x320 grid. Having these large wires between the MAC units may cause further complications when it comes to bringing the hardware into reality, whether in FPGA or an actual tapeout (physical silicon).

Our original design is closer to what modern ASICs⁸ actually use for accelerating matmuls, and that design is famously known as the Systolic Array!⁹

Over here we can see the differences in the wiring of our original design and the design we settled on.

Systolic array1 / 2

A goal we had for this project was to build an intuition for invention, rather than building pre-existing architectures; therefore, we went ahead with our version of the MXM!

MXM1 / 6

09

SXM

Now looking at that matmul, there was a strange T above the K vector.

That “T” denotes something called a Transpose. When neural networks are performing calculations using the Attention formula, they usually need to perform this Matrix Transposition. Simply put, transposing a matrix means the orientations of the rows and columns get flipped diagonally. If we didn’t transpose the matrix, computing Attention would be extremely hard since the MEM modules and MXM can only read and multiply in parallel, row by row. For example, if the Q and K matrices are being multiplied, the data in Q would be read row by row while K would be read in columns. This would result in one piece of data being read per clock cycle, completely stalling the pipeline.

The critical path is only bound by the speed of the matmul in the MXM, not transpose in the SXM. Transposing a matrix would generally cost extra clock cycles, but since our design is statically scheduled, there is no extra cost. The compiler schedules the instructions for the SXM to rearrange or transpose the data, and it is done while the data is travelling to the MXM!

10

Activation Functions and Neural Networks

Now that we have discussed the need for matrix multiplication and how it works in hardware, we should talk about the next most important thing in ML: neural networks!

In an actual brain, neurons receive electrical signals from their surrounding neurons, and when it reaches a threshold, the signals are passed on to the next neuron. These “AI” neurons work in the same way; they take in data, perform calculations, and decide whether or not to pass that information onto the next layer.

Above is an image of a simple neural network. In every line of the network, the formula y = mx+b is used. As we move down the network, we keep stringing that same formula together to end up with a straight line at the end of the network.

When data enters a neuron, there are three components that are calculated. “x” is the input data, “m” is the weight (which is how much importance the input has), and “b” is the bias, which is the shifting of the line up or down. These neurons are stacked on top of each other to create a network.

Let’s walk through an example:

Y = 4x+7 is still a linear equation. You could stack an infinite number of neurons on top of one another, but the output will always stay linear. This becomes a major issue as some data just can't be classified with a straight line.

Let’s take a look at this graph to understand how it really works. You want to classify two fruits, lemons and tomatoes, on two different features, redness and tanginess.

The data on the graph forms groups. The tomatoes are red but not as tangy, so they group up at the bottom right corner, while lemons are not red and very tangy, so they group up on the top left. The line in between the two groups is the neuron. It uses y = mx + b to create a boundary, so everything on one side is a part of the tomato group, and everything on the other side is a part of the lemon group.

Now imagine a trickier scenario where we want to classify oranges and lemons. Oranges are quite tangy, but they sit closer to redness than lemons do. The visualization of this data places some orange data points trapped directly inside a surrounding cluster of lemons.

This is what non-linear data is. Imagine trying to draw a line to split all the lemons and oranges in 1 go with no overlaps on either side; you can’t. No matter how you orient that line, you can't split them without either going through the oranges in the middle or splitting them unequally.

As we saw earlier, simply stacking more linear layers together just collapses back into another straight line. To solve this issue, we don’t throw away y = mx+b. Instead, we actually blend the output of each neuron with a non-linear function known as an activation function. The simplest and most popular of these activation functions is known as ReLU (Rectified Linear Unit).

ReLU is a concept that runs on a very simple premise:

If a number is negative, set the output to 0.

If the number is positive, let it pass through.

By following these rules, ReLU breaks the linearity. ReLU will take that line of yours and start to bend it. By combining these bent lines together, we start to create a boundary around the lemons to group them together, leaving the oranges in the center untouched.

After adding ReLU to our graph, we can now see that there is 1 closed loop line that creates a full boundary around the lemons and oranges.^{10, 11}

There are also other activation functions such as Sigmoid and GeLU / SiLU / SwiGLU¹²; however, ReLU provides the perfect balance between simplicity and effectiveness, so we decided on using that.

Now that we have a background on neural networks and activation functions, we can dive into how we represent this in hardware with our VXM.

11

VXM

VXM v1

The Vector Execution Module (VXM) is the module responsible for scalar and vector operations.

The LPU was designed with the compiler in mind, and Groq believed that with basic operations such as addition, subtraction, multiplication, exponentiation, etc, the compiler could string these operations together to form complete formulas. When we started work on the VXM, we wanted to get the basics down, so we started with 3 operations:

1. Bias Add (addition)
2. Scale (multiplication)
3. ReLU

With these 3 operations, we wanted to make formula creation (stringing operations together) easy for the compiler, so we decided that adding a selection of which operation you want at a time would be best. To do that, we needed to use a hardware block called a multiplexer (MUX).

A MUX is super useful and can almost be visualized as a train-track selector: if I flip a switch to the left, it should choose the left lane; if I move the switch to the middle, it should pick the middle lane, and so on. Instead of a switch, they work on control bits. If I pass a zero into the control line, lane 0 is my output; if I pass 1, lane 1 is my output. This can be scaled to as many lanes as you need, but for our purposes we would require a 3-to-1 MUX, where the control bit is now 2 bits wide instead of 1 bit.

When the compiler sends its instructions, 2 of those bits are used to control our 3-to-1 MUX inside our VXM and choose which of our 3 operations to output. Thus making our first version of the VXM:

However, this design didn’t come without flaws; its main flaw was that data would have to completely exit the VXM before coming back and hitting another operation. On a bigger scale with more operations, this would become a major traffic issue.

Another flaw we noticed was that this only operated on one element at a time, whereas a row of matmul would have multiple elements at a time.

For example, if we had to run add on 2 data lines and needed to run scale after, we would have to completely exit the VXM and feed back into scale. For a project that only needs bias addition and an activation function, this would be more plausible. However, for our Language Model we would need softmax, RMSNorm, Vector Add, etc, which forced us to go back to the drawing board and redesign this architecture.

VXM v2

When we redrew this architecture, the main problem we had in mind was “how do we implement the math so that the element never leaves the VXM before coming back in?”

So we had the idea of “splitting” up the operations into “phases.” Phases being operations such as bias add, ReLU, Scale, etc.

VXM v2 pipeline1 / 4

So while an element is in ReLU, the next queued-up element can enter a bias addition. This creates overlap between operations! As the animation above shows, instead of forcing element \(a_{01}\) to wait until \(a_{00}\) finishes both bias addition and ReLU, we decouple the two stages. While a00 is in ReLU, a01 can start Bias addition. This concept is known as throughput, and with our pipeline we have increased throughput as more operations are starting at an arbitrary point in time compared to v1!

Now sometimes we may not need to do certain operations. For example, when the result of \(Q @ K^{\mathsf T}\) needs to get scaled by \(\sqrt{d_k}\), we don’t need that value to go through bias add or ReLU.

That is when we decided to add a MUX between each stage! By adding a MUX, we can include the wire before it goes through an operation, allowing us to bypass it based on the select bits sent via the compiler. This is when we truly started to understand how much power the compiler truly holds for the LPU.

To address the flaw of operating on multiple elements at a time, we expanded the VXM to hold multiple lanes. For example, let's say we have a row [a00, a01]; as of yet, we can only do computations on one element at a time. We wanted to parallelize this process by being able to execute an operation on each row at a time. That is when we introduced multiple lanes!

After all of the elementwise operations (bias add, ReLU, Scale) are completed, we collect them into a row to proceed with the vector operations.

Now it's time to implement vector-wise operations, including Softmax, RoPE, RMSNorm, and Vector Add¹³.

Full VXM v21 / 13

Groq designed their LPU to be extremely flexible. Inside their VXM, they didn’t have a physical softmax operation. Instead, they designed their hardware with basic calculators that can do operations like add, multiply, finding maximums and minimums, etc. To execute operations like Softmax, they completely rely on the compiler. The way it works is that the compiler breaks down the complex request into tiny ones that the hardware can perform using simple arithmetic blocks. For example, if they wanted to calculate:

The compiler tells the hardware to send data x and a into a multiplier block, take that result and send it into an adder block to add b, then take that result to send it into another multiplier block to multiply with c. This design choice keeps the silicon flexible; if a new activation function is discovered, only the compiler needs an update, not the physical chip. However, making a compiler smart enough to schedule hardware at this micro-level takes massive engineering effort, which is why big tech spends millions to recruit elite compiler engineers. Plus, all those instructions take up lots of memory and bandwidth.

Since we are not a big tech company, we took the opposite approach. Our focus was on hardware, so we didn’t mind sacrificing flexibility as long as we maintained some sort of determinism within the chip. Hence, we hardwired RMSNorm, Softmax, and RoPE directly. While we sacrificed the ability to make quick updates if there are any big changes, we have the massive advantage of speed.

12

Look Up Tables

Implementing arithmetic such as multiplication and addition on hardware thus far was really straightforward; however, our first real bottlenecks when it came to pure arithmetic appeared during the implementation of softmax, where we had to implement division and exponentiation (as seen in the formula).

Implementing division and exponentiation led us to many stalls, as together they took over hundreds of cycles to finish computing. Multipliers and adders are relatively easy to build out of logic gates, making them cheaper operations, but one 8-bit division module requires roughly 200 to 1,500 logic gates! All of those gates add computation time; repeat that operation over and over, and our LPU would take minutes for just a single token output.

But what if the hardware knew the answers to division/exp questions as they were asked? That is when we learned about Look Up Tables. A Look Up Table (LUT) is a table that stores pre-calculated output values for a set of inputs. Instead of performing a calculation every time, the system simply looks up the corresponding value in the table, which is much faster. This also has an added feature to approximate the result of an operation. Now the term “approximation” seems very scary for hardware, or at least it did to us before we dove deep into it, but look at it like this:

Let's look at the example \(\frac{10}{3.107}\)

Division can also be interpreted as multiplying the reciprocal

Therefore, \(10 \times \frac{1}{3.107}\)

This circles back to the idea of arithmetic in hardware; multiplication and addition are cheap. Now technically \(\frac{1}{3.107}\) involves division, but that is where LUTs come in!

We store a bunch of reciprocal values into memory such as:

Now let's say we have to do \(\frac{10}{3.107}\).

In reality, the answer to that is \(3.21854\)

In hardware using LUTs, if we do \(10 \times \frac{1}{3.107}\), we look at the closest value to \(3.107\), which comes out to \(3.11\)

The hardware now “looks up” in the table to see what the value of \(\frac{1}{3.11}\) is, which comes out to \(0.32154\)

Now we do the final arithmetic \(10 \times 0.32154\), and get \(3.2154\)

Now, if we compare the values

Even though we approximated the value of this division, we got really close to the expected value. We essentially turn the division arithmetic into multiplication. The tradeoff in this comes out to be that there is a small error when comparing the actual value and approximated value.

The error value is hard to be zero; however, it can be as close to zero with design decisions. For example, if you see the divisor column in the LUT above, it is spaced out very lightly, incrementing over \(0.01\) every entry. However, if we make it less spaced out with something like:

Now the divisor increments by \(0.25\); the closest value to \(3.107\) is \(3.00\), and in that case \(10 \times 0.3333 = 3.333\)

Which comes out to an error of \(3.33\%\)

An LUT with divisor values in very small increments is more accurate, but takes more space in memory; however, an LUT with larger increments is less accurate, but takes up less space in memory.

13

Buffering

During an F1 race, one of the most crucial parts is the pit stop. The fastest pit stop ever is 1.80 seconds, meaning the pit crew took off 4 tires, acquired the 4 new tires, and then screwed them back on all under 2 seconds.

Now imagine if the pit crew had to look for the newer tires after unscrewing the used ones. They would have to unscrew the tires, look for the newer ones, carry them to the car, and then screw them back on. Achieving a time like 1.80 seconds would be impossible under those conditions. Therefore, the pit crew finds the tires they need before the newer tires are even unscrewed.

This concept is known as Buffering, where, in simple terms, you store the next elements you need for a computation while a current computation is taking place. Let's look at the MXM; currently, once the MXM is finished computing a dot product on current inputs, we have to wait until the next inputs get loaded before the next computations can start. To mitigate this issue, we introduced another buffer on top of the one we already have, so while the first dot product is being computed, we can preload the next elements!

Before this buffering, the MXM took 32 clock cycles to compute an 8x8 @ 8x8 matmul. After introducing double buffering, we were able to decrease that count to 25 clock cycles, approximately a 21.875% reduction in total timing!

MXM double buffering1 / 5

This same principle can be applied for outputs! Let's say we need to start a new matmul, but the VXM is not ready to consume the current output. In that case, we can add registers on the output side to hold the current computation, allowing us to start the next computation.

We also applied buffering inside the VXM. We have 4 input buffers per lane, and 4 output buffers per lane.

14

RMSNorm

Throughout this article, there's been a lot of complex math mentioned (and more to be mentioned); this chaotic nature of transformers can cause values to either shrink to zero or grow too large for the model to handle. To mitigate this phenomenon, normalization is used. Normalization is like an automatic volume controller, it makes all the audio comfortable to listen to whether it’s a whisper that gets louder making it easier to hear, or an explosion that gets softer. Traditionally, LayerNorm is used, but RMSNorm provides a simpler formula with a similar effectiveness, so we went ahead with it. The RMSNorm formula can be represented as:

Our initial bottleneck came with the complexity of the formula, as it includes

1. The addition of \(\varepsilon\)
2. Square root
3. Division (\(\frac{1}{d}\))
4. Finding the reciprocal (later in the calculations)

Which are all arithmetic operations that in hardware are generally tricky to do. That is when we applied the trick we learned earlier: Look-Up Tables!

Now the first problem is very easy to deal with: we simply can just remove \(\varepsilon\)! The reason why we can remove this value is that it was used as a failsafe to prevent a division by zero. Now that seems important, but as discussed before, to solve the division issues we can just use an LUT and have a lookup value that, when we see \(\frac{1}{0}\), we set it equal to a value as if it were \(\frac{1}{0.001}\).

Before we thought about LUTs, we observed how a simple 4 row of RMSNorm is computed to notice any patterns.

Seeing this calculation was very important! As we noticed, we reciprocate \(\sqrt{x}\). Therefore, instead of using an LUT for first calculating the square root, and then the reciprocal. We combined them into one, where we looked up the value of \(\frac{1}{\sqrt{x}}\).

As for division, we initially thought we would need an LUT for division as well for computing \(\frac{1}{d} \times \sum(x^2)\). But then we realized any divisor which is a power of 2 can simply be right shifted to perform division!

Now right shifting bits is a very powerful operation in hardware, and it can be used to replace division!

Let's look at how bit shifting works in our RMSNorm:

Let's start with our reciprocal \(\frac{1}{d}\). In RMSNorm, d is equal to the number of features in your model. In our case, since we are using MicroGPT, our feature size is 16, so \(d = 16\).

Now binary works on a power-of-2 system. Take this 8-bit binary number for example.

The first bit (the bit on the far right) is in position 0, so we raise 2 to the power of 0. For the second bit, it is in position 1, so we raise 2 to the power of 1, and so on.

But remember how our d value was equal to 16; 16 is just the same as \(2^4\). The general rule for right shifting is that any division by a power of two (\(2^k\)) is equivalent to a right bit-shift by k. Meaning that since our \(k = 4\) to divide by \(16\) (multiplying by \(\frac{1}{16}\)), it is the same as right shifting by 4 bits.

To complete the division, we will then take our binary number and physically shift the bits. Here is an example.

Take this binary number, which equals 80.

Then apply a right shift of 4 bits (\(\frac{80}{16}\))

The resulting value is a binary number which is equal to 5! (\(\frac{80}{16} = 5\))

Now the problem arises when the number is not a power of 2, such as 6. That is when you would need an LUT. However, for the model we executed, the trick remains true and saves us lots of memory bandwidth!

This is the first time we battled with the tradeoffs of LUTs; at first, creating 3 different LUTs would be the easy thing to do. The tradeoff being that LUTs occupy memory, and if you want an LUT with a small error% (as shown here), the amount of memory needed could be very large.

In this instance, we got creative and only used 1 LUT and had some leeway to make it with a small error%.

After the implementation of LUTs, we ran into another bottleneck. We noticed that there is a summation across the entire row. For example, if we had a row of 16 elements, all 16 of these elements must be accumulated before the operation can continue. Working with our 8-lane VXM, we thought of a different approach: what if we process the first 8 elements, and store their partial sum? Then, when the next 8 elements arrive, what if we just accumulate those with the values we already have? This would progressively build the final row’s sum rather than have us expand to 16 lanes.

After those 16 elements have been processed, we can then move to the next stage and square them, followed by multiplying by \(\frac{1}{d}\), and then square rooting them.

This idea is known as chunking!

In this animation, you can observe how we used Chunking, and LUTs to compute RMSNorm in hardware.

RMSNorm1 / 20

Chunking at its core felt similar to pipelining for us, but on an internal level. Instead of computing RMSNorm in one go, we split up the operation into 2! Not only did this allow us to process a larger row than what the current size of the VXM is, but the current row that is being processed does not cause a stall for multiple clock cycles. For this example, we can allow for 2 rows to enter for one computation.

15

Softmax

Remember when we introduced activation functions? Well, it's time to dive into another one; this time into a function called “softmax.”

Softmax is used to turn an array of random numbers into a set of probabilities all adding up to 1.0.

Remember when we mentioned that doing \(Q @ K^{\mathsf T}\) gives a similarity score for how similar the vectors are? We use softmax to give percentages of context!

For example, let's have \(Q_1\) represent the vector for “thinking”

\(K_1\) represent the Key vector for “thinking”

\(K_2\) represent the Key vector for “machines”

Without context, “thinking” and “machines” are separate.

When we do

We can infer from these results that, for the word "thinking" in this sentence, approximately 80% of the information comes from the word "thinking" itself and 20% comes from the word "machines." In other words, when the model looks at “thinking,” it also pays attention to “machines,” helping it understand that the sentence is discussing machines that think.

Softmax is represented by this formula:

Our friend Surya had actually already implemented softmax in hardware before¹⁴. It included a divider module, exponential module, and the means to chain them all together, so we just ported this module into our own VXM. When we eventually tested the functionality, we noticed that this softmax gave correct answers but took around 300 clock cycles to complete.

Therefore, in order to compute the division and exp arithmetic here, we anticipated that we had to use both LUTs and chunking.

Previously in RMSNorm, we omitted the division LUT due to the bit shift trick; we looked to do something similar in softmax, but were unable to find any workarounds. Within the denominator, you sum a bunch of exponentiated numbers. As we recall, the bit shift trick only works if the denominator is within a power of 2.

Chunking was necessary as softmax distributes probabilities over a row, meaning if you have a row of 8, but your VXM can only process rows of 4 at a time, you will have to split these two rows, but computing softmax separately on these rows would give 2 rows having their own probability distribution, meaning if you add them both the sum would be 2, when we want the sum of both rows to be 1. The following animation shows how the softmax module computes on our chip!

Softmax1 / 27

16

RoPE

As previously learned, we can represent this sentence as vectors.

However, one flaw at the moment within the embeddings is that the model has no idea which order these words are in. If we extend this sentence to something like “thinking helps machines.” The model may understand which action is being done, but it may not be able to distinguish between thinking, help, and machine, as it is not aware of their positions. “Machines help thinking” is mathematically similar to “thinking helps machines.” So here comes RoPE (Rotary Positional Embeddings)!

The entire idea for RoPE is to “rotate” a vector based on the position.

In this graph, you see vectors for 5 tokens. The first one (m = 1) is at the vector's original position; m = 2 is rotated slightly away from it, m = 3 is rotated even more, and so on. So when the model observes these tokens, it notices immediately which one came first, second, third, etc. Furthermore, none of the actual information the vector wants to convey gets altered; only the direction of the vector changes. This is a way for the model to see where the word is in the sentence geometrically in space without remembering or counting¹⁵.

We implemented three separate LUTs directly in the hardware:

1. The frequency for the angle calculation (\(1000^{-\frac{2}{d_{\text{model}}}}\))
2. Sine calculations
3. Cosine calculations

Unlike previous modules such as RMSNorm and Softmax, we did not need multiple passes of chunking here. For the previous computations, we had multiple passes as they required data for an entire row at a time to provide precise results, whereas in RoPE computations are performed in pairs. Each pair is independent and has different numbers. The process of RoPE being computed on our LPU can be visualized with the following animation.

RoPE1 / 19

17

Putting the math together

Now that we have all our mathematical blocks ready, we can have a look at the self-attention formula one last time:

So far we have learned everything up until \(Q@K^{\mathsf T}\) and softmax; however, the only things that remain are \(\frac{1}{\sqrt{d_k}}\) and \(@V\).

The division by \(\sqrt{d_k}\) is used to prevent the values of \(Q@K^{\mathsf T}\) from becoming too large for softmax.

For example lets say we have a row [20, 30, 40, 50]. If we compute softmax over that row, we get [0, 0, 0.000045, 0.99995], whereas a row of [2, 3, 4, 5] may give us [0.03, 0.09, 0.24, 0.64]. The latter is much more feasible due to what follows next: multiplying with V.

As we learned prior, the V vector is responsible for holding the actual contents for that specific word. So let's look back at our example of “thinking machines.”

Let's say for the word “thinking” we get softmax results of [0.8, 0.2]. 0.8 being thinking, 0.2 being machines.

We multiply 0.8 by the V vector of thinking, and 0.2 by the V vector of machines.

After that, we combine those two vectors into one. Now we have 80% of the context from thinking, and 20% from machines!

Congratulations! Now you understand how Language Models give context to each word!

18

Quantization

When we went to store these values into MEM, we realized two things. When operating on the MXM values, they accumulated them into int32, and the VXM operated on int32. However, we wanted to store these as int8. The reason being that while larger data types hold more accuracy and features, they also become more expensive to store in hardware!

To compress these values down to int8, we can perform quantization!

Quantization must be done with precision, as with int32 you have 4.2 billion unique values, and int8 you only have 256 unique values. Therefore, we treat each row uniquely.

To expand,

We can have one row with extremely large values like

[1200, 1800, 2400, 3000] in int32

And another row like

[200,300,400,500] in int32

If we treat both rows the same and compress them down to int8, the latter row can become extremely small! We can denote the quantization of these two rows as:

After these calculations, we can notice that quantization is actually really simple! You can follow the steps.

1. Find the max of the row
2. Divide that max by 127 to get the scale factor
3. Divide the rest of the row by that factor

As you can see, the highest value will always be 127. Even though the two rows originally had different maximum values, each row requires a different scale factor to map its values into the int8 range. We hold the scale factor in MEM so it can be reused later to multiply back to its original values, or as close as possible.

Quantization is important for the LPU because the LPU only utilizes SRAM, which is very fast to read from, but does not have a lot of storage. Due to the scarcity of storage, quantization allows the LPU to better manage its memory. For example, the current size of our LPU memory comes out to be ~288KiB; this is for everything in int8 in memory. For 32 bits stored in memory, we would need ~1,056KiB. That is around 3.67x the memory we currently have!

19

ICU

Now let's discuss how the compiler actually talks to the LPU. In the LPU, we have something called an Instruction Control Unit (ICU). As we have mentioned, the core idea of the LPU is that the hardware is dumb and doesn't know how to handle different types of AI models; the compiler is what tells the chip how to run them. The compiler generates a 96-bit-long binary string and sends it to the ICU. When the ICU receives this string, it starts to unpack it and send specific bits off to different parts of the chip.

Each bit in the string is assigned to a different location on the chip, and it looks like this,

We've already discussed the parts of the chip these bits go to: the MUXs! The select bits in each MUX are fed directly from the ICU and lead to the deterministic approach of the LPU.

Let's take a quick look back at the VXM for a second.

You can see how the output of scale and the input of scale both feed into the MUX. The compiler will send its instructions to the ICU, and then the bit assigned to that MUX decides whether scale should be used or bypassed, all happening in nanoseconds!

The ISA (Instruction Set Architecture) handles data movement very efficiently using a system that we designed. The first path is the westbound bus, where data can only travel from the east to the west. The second path is the eastbound bus, which handles data movement from the west to the east. To keep the system flowing fast without any stalls, each direction only has one unit sending data and one unit receiving it per instruction. However, both the east and westbound buses can be used together for a single instruction. We built the buses using a simple MUX-to-DEMUX¹⁶ architecture. It works like a railroad switch where the compiler is the one controlling the switch. A combination of bits goes into the MUX, which flows from one of the buses, and is received at the other end through a DEMUX. From there, it sends the data to the correct module.

20

Memory Hierarchy

The memory in an LPU is stored on the chip itself in something called SRAM, which is just a collection of buffers and registers. This is very different compared to the memory hierarchy of other AI accelerators.

This former diagram is the memory hierarchy in a GPU¹⁷. We notice a relationship with the size of storage: the smaller the storage, the faster it is to retrieve an element; the larger the storage, the slower it is. Registers and Cache cost more per area on the chip, and the slower memory storage, RAM and SSD, are cheaper. When they are all used together, the cost is minimized, and performance is maximized.

The LPU decided to completely go against memory hierarchies so it does not have to keep fetching data and instructions from DRAM or HBM¹⁸. Instead, the only MEM that exists on the LPU is on-chip memory, meaning the memory you see in the animation is the only memory that exists.

While working with transformers, we run into two specific bottlenecks: memory-bound workloads and compute-bound workloads.

When the processor spends more time retrieving elements from memory than actual arithmetic operations, a memory-bound workload occurs.

A compute-bound workload is when the arithmetic operations are the bottlenecks. The memory feeds the arithmetic units, but they are not fast enough or have reached maximum capacity on the task at hand.

These two types of bottlenecks appear during different stages of text generation. There are two main parts that take place when we generate text: prefill and decode.

During prefill, the entire prompt is read all at once, so let’s say the goal is to output the name “Aaron”; we would input Aar, which would be read and processed by performing matmuls.

When this is happening, the arithmetic units are at max capacity, making this part compute-bound.

Decode is when the response is being generated. To generate an output, the chip loads a massive number of parameters from the memory, does some arithmetic, and repeats the process. This part is memory-bound because the arithmetic units aren’t being used as much while the memory units are working at max capacity.

To solve both bottlenecks, we use something called a KV cache; the KV stands for Key-Value. This is the chip’s short-term memory. During prefill, the keys and Values of previous tokens are saved in a cache. This prevents unnecessary calculations and increases the speed of the output during the decode phase.

However, since an LPU does not follow a memory hierarchy and only includes SRAM, this means it does not have a large amount of storage, making the KV cache harder to maintain, whereas DRAM or HBM have large amounts of storage capabilities, making it a much better fit for KV cache.

Coming back to how memory on an LPU works, the entire hierarchy is completely disregarded!

Large amounts of SRAM (caches) are on the chip, so since there is no hierarchy to follow, the LPU does really well navigating memory-bound workloads like decode.

Let’s take a look at traditional GPUs really quickly. GPUs are packed with thousands of cores that perform arithmetic calculations, so they are amazing at compute-bound workloads like prefill. But they rely on the memory hierarchy, so they aren’t as great at memory-bound workloads like decode.

After learning this, we guessed that Nvidia acquired licensing rights to Groq’s LPU to build some sort of inference disaggregation system that does compute-bound workloads on the GPU, and memory-bound workloads on the LPU.¹⁹ This was later proved correct when Nvidia CEO Jensen Huang revealed the reason for the acquisition at GTC²⁰. This also mitigates potential KV caching issues an LPU would hold by using the GPU for the KV cache!

21

Numerics

To represent the actual numbers that the chip used, we originally used an 8-bit integer for our data. Integer arithmetic is the easiest to implement in hardware. It is also the cheapest and fastest when it comes to multipliers and adders relative to other formats. Inputs, Weights, Outputs, or anything that is to be stored in the MEM blocks is 8 bits. 8 bits was the perfect compromise between speed, efficiency, and numerical accuracy.

However, the nature of int8 did cause a couple of problems. Signed Int8 is in the range -128 < x < 127, and in traditional int8, there are no decimals. Meaning the range of values was super small (only 256 values in total) to represent individual features.

We then researched numerical formats used for modern ML accelerators, and we came across the two most popular formats, floating-point & integer-fixed-point.

Floating-point-8 works like this:

Therefore, floating-point math is really convenient for machine learning as it has a large numerical range! However, for how convenient it is, there are also many drawbacks. For starters, in 8-bit integers, a common multiply or add might be almost instant, whereas for floating-point it has to factor in signs, exponents, etc before multiply or add can be done. This makes the complexity of floating-point very high.

This helped us realize that for this project a numerical format as complex as floating-point may be too “overkill.” We went on with it anyway, using the CVFPU library²¹. Upon our first step at synthesizing this design, we ran into many errors as this library was incompatible with our FPGA (despite working well on testbench simulations²²).

We then pivoted to fixed-point. Fixed-point numbers work differently; imagine having to represent the amount $12.34 in cents. You get 1234; however, you know that the cents will always be the 2 spots after a decimal. That is the motivation behind fixed points; the word “fixed” comes from the fact that there is a fixed place where the decimal value will be.

Though fixed points do have the benefits of providing a larger range of numerical explanations, it also mitigates the complexity of hardware, which made this a home run decision for us.

22

FPGA Implementation

Note: This info is a little harder to grasp; it is for those who are interested in how we synthesized the LPU on an FPGA.

We used a Terasic DE1-SoC development board to implement our design. We used CoCoTB to test all the modules and our logic during the design process. Testbenches were really important as they helped us identify any mistakes and test edge cases for our RTL designs (we were able to identify the softmax division bottleneck because of this!) After the completion of our design, we measured ~8.5k tok/s (tokens per second) when running inference of MicroGPT using CocoTB.

When we brought up the design on the FPGA, it was extremely dense and ended up using 100% of the potential logic blocks, which forced us to lower our clock frequency from 100 MHz to 6.75 MHz. This means we would have a clock edge every 160ns on the FPGA, from 10ns on CocoTB. This disparity was one of the reasons why our FPGA tok/s came out to be 20. Theoretically, the 8500 tok/s can be brought down to 530 tok/s simply due to the clock frequency decrease; the rest of the reduction of tok/s can be attributed to other overheads.

23

Final Thoughts

Looking at our journey, we are extremely proud of what we have built after starting off with no prior chip design knowledge. Our biggest takeaway from the project was to approach everything with a first-principles approach. You don’t have to know everything before trying something crazy. You can learn as you go. By breaking the project down into tiny pieces, we were able to solve problems that seemed impossible at first glance.

Our biggest weapon was not some crazy tool that nobody knows about; it was a pen and a piece of paper. Drawing everything out, no matter what we were working on, helped us visualize the problem at hand and helped us immensely in the ideation process. It made us think about the hardware physically before a single line of code was written.

As we had mentioned at the beginning, one big reason we wanted to work on this project together was to move away from being fully reliant on AI. We drew everything out, discussed different ways to approach problems and solutions to those problems, and only after that did we use AI to accelerate the coding after understanding every element of our design. We had complete ownership and proved to ourselves that in today’s age, we could still think and figure things out ourselves.

You might be wondering, what’s next?

Our current focus is building a compiler to handle all the instruction scheduling, tiling and chunking and have the LPU run completely on our custom instructions.

Since we have proven that our design works on a small parameter model, our next goal is to push the hardware to run on a much larger model, from 4.9k to 260k and scale our memory and compiler along with it.