Understanding llama.cpp — Computation Graph and Transformer Architecture

In the world of Large Language Models (LLMs), the Transformer architecture gets all the attention. But behind the scenes, frameworks like LLaMA rely on a sophisticated graph-based computation engine to handle memory, optimize execution, and achieve blazing-fast inference.

In Part 1: Understanding llama.cpp — Efficient Model Loading and Performance Optimization, we explored how LLaMA efficiently loads and optimizes large models. Now, in Part 2, we take a deep dive into how LLaMA builds, optimizes, and executes computation graphs.

What You'll Learn

• ?? Understanding LLaMA's Transformer Architecture: Layers Unveiled
• ?? Why LLaMA Uses Graph Computation Instead of Direct Calculation
• ?? Anatomy of LLaMA's Computation Graph: Nodes, Edges, and Tensors
• ??? Building the Graph: How ggml_build_forward_expand Assembles the Transformer
• ?? Topological Sorting: Decoding the Execution Order of the Graph
• ?? Running the Graph: How Inference is Performed Step-by-Step
• ?? Key Takeaways: Why Graph-Based Computation Makes LLaMA So Efficient

If you’re aiming to understand LLaMA's inner workings or build your own LLM engine, this article is for you.

?? Understanding LLaMA's Transformer Architecture: Layers Unveiled

At its core, LLaMA follows the Transformer architecture, but with optimizations to reduce memory consumption and increase execution speed. Here’s a simplified view of the key components of the Transformer used in LLaMA:

1. Input Embedding: Converts input tokens into continuous embeddings.
2. Positional Encoding: Adds positional information to the embeddings.
3. Multi-Head Self-Attention: Allows the model to focus on different parts of the sequence.
4. Feed-Forward Layer: Applies non-linear transformations to the output of the attention mechanism.
5. Normalization: Used to stabilize training and inference.

LLaMA's Transformer layers are stacked N times (where N is the depth of the model). Each of these layers has several intermediate tensors that must be computed, stored, and reused. Directly computing these tensors at runtime would be inefficient, which is why graph-based computation is introduced.

?? Why LLaMA Uses Graph Computation Instead of Direct Calculation

In a traditional implementation, you compute tensor-by-tensor, step-by-step. For example:

1. Compute tensor A.
2. Use A to compute B.
3. Use B to compute C.

This method works but requires memory to hold every intermediate result, and it prevents global optimizations.

In graph-based computation, you first define the entire graph of operations and tensors. Each node represents a tensor or an operation (like addition, multiplication, etc.), and edges represent the dependencies between them. Once the graph is defined, it is executed using an optimal order (using topological sorting).

Why Use Graph-Based Computation?

• Efficient Scheduling: Operations can be parallelized.
• Memory Optimization: Intermediate results are not stored indefinitely.
• Lazy Execution: Computation only happens when required.

LLaMA leverages this by using the ggml library, which allows you to build the graph of computations using simple functions like ggml_add, ggml_mul, etc.

?? Anatomy of LLaMA's Computation Graph: Nodes, Edges, and Tensors

LLaMA relies on two key C++ data structures:

1. ggml_tensor: Represents each operation and its data.
2. ggml_cgraph: Represents the entire computation graph, including nodes, leaves, and execution order.

Data Structures (Simplified View)

??? Building the Graph: How ggml_build_forward_expand Assembles the Transformer

Now, let's walk through how LLaMA builds its computation graph for a single Transformer block.

Step 1: Inputs and Embeddings

The first step is embedding the input tokens. These embeddings are stored as leaf nodes in the graph since they’re not recomputed.

Step 2: Self-Attention Graph

For each input, we compute the queries, keys, and values. Each step (like Q = Embedding * Wq) becomes a node in the graph.

Each function like ggml_matmul, ggml_transpose, etc., does not compute anything immediately. It simply adds a node to the graph.

Step 3: Feedforward Network

Output from attention is fed to the feedforward network. Similar to attention, the intermediate calculations are stored as graph nodes.

?? Topological Sorting: Decoding the Execution Order of the Graph

Once the graph is built, we need to determine the order of execution. This is achieved using a depth-first search (DFS) to topologically sort the nodes.

How LLaMA Sorts the Graph

?? Running the Graph: How Inference is Performed Step-by-Step

Finally, we execute the graph node by node using the topological order.

How LLaMA Executes the Graph

?? Key Takeaways: Why Graph-Based Computation Makes LLaMA So Efficient

• Graph-based computation avoids unnecessary recomputation.
• Lazy execution optimizes memory, freeing tensors when they're no longer needed.
• Topological sorting ensures execution order respects dependencies.