Luminal Internal Representation Explained

Going Through How Luminal Represents and CodeGens Naive Matmul

Image by Author — Flux.1

Since we figured out how Luminal can discover Flash Attention, we’ve been getting a lot of questions about how this was done. In short, we created our own internal representation (IR) that describes machine learning models and used that to search the space of equivalent kernels. To understand how this is done, it’s helpful to go through an example. Flash Attention itself is quite complex, so we’re going to focus on the simple example of naive matrix multiplication (naive matmul).

If you’d like to see exactly how Luminal discovers Flash Attention, check out the Github here, or to see the intuition behind Flash Attention check out my blog here.

What You Need to Know Before We Start

Let’s start with the basics. Luminal uses a tool called egglog to explore equivalent ways to compute the same math. Think of it like this:

Internal Representation (IR): A simplified “blueprint” of how your model computes values. It’s not code — it’s a math-y abstraction that makes optimization easier. e-classes: Groups of operations that behave the same way. For example, 2*2 and 1+1+1+1 are different e-classes but produce the same result. Egglog searches these groups to find the fastest version. Flash Attention: A clever algorithm that speeds up attention layers in models like Transformers. Luminal can “discover” it automatically by searching the space of equivalent computations using its IR.

We’ll focus on naive matrix multiplication (matmul) to show how this works, but the same tools let Luminal find Flash Attention. Cool? Let’s go!

E-Classes

Luminal uses egglog to search the space of equivalent kernels. To use egglog, we define operations, any mathematical properties, and ways to transform one operation into another (think how we can either perform 2*2 or 1+1+1+1). All of these rules are written out in a LISP-style.

Why does egglog care about equivalent operations? Because some versions of the same math run faster on GPUs. For instance, 1+1+1+1 might be slower than 2*2 due to instruction pipelining.

For naive matmul, we only need 6 e-classes: Mul, Add, LoopIn, LoopOut, MAccum, and Tensor. Tensor simply defines the data in our IR, so we will largely ignore its instantiation here. Add is an element-wise addition for a tensor, and Mul is an element-wise multiplication. LoopIn and LoopOut tells Luminal what code is meant to run on a loop. The first argument in to LoopIn holds the size of the data and the second is the stride we should use when accessing it. All these operations run on Tensors, so it is basically a given that LoopIn and LoopOut will be involved somehow. Finally, MAccum accumulate values like summing results in a dot product.

Element Wise Multiplication

To show how this comes together, let’s do a simple element-wise multiplication. We’ll take 2 Tensors (A and B) that have the same shape (m x n). Then we write out our IR like so:

(LoopOut (LoopOut ( 
       (Mul (LoopIn (LoopIn A m n) n 1) 
             (LoopIn (LoopIn B m n) n 1)) 
n 1) m n)

That code looks like this in C:

void element_wise_mul(double* A, double* B, double* C, int m, int n) {
    for (int i = 0; i < m; ++i) {
        for (int j = 0; j < n; ++j) {
            C[i * m + j] = A[i * m + j] * B[i * m + j];
        }
    }
}

Naive Matrix Multiplication

Because matmuls are not element-wise, we need to do some broadcasting to ensure that even though we use Mul, we are still getting the expected result of a matmul at the end.

Image by Author — Matrix Multiplication Example (not element wise)

To handle this broadcast, we are going to add in a third loop of stride 0. By doing this, we give ourselves a new degree of freedom with which to access the second matrix without adjusting our place in the first matrix. In our IR this looks like this:

(LoopOut (LoopOut (LoopOut 
    (Add 
        (Mul (LoopIn (LoopIn (LoopIn A m k) n 0) k 1) 
             (LoopIn (LoopIn (LoopIn B m 0) n 1) k n)) 
        (LoopIn Acc k Acc)) k Acc) 
n 1) m n)

In C this looks like this:

void naive_matmul(double* A, double* B, double* C, int m, int n, int p) {
    for (int i = 0; i < n; ++i) {
        for (int j = 0; j < p; ++j) {
            double sum = 0.0;
            for (int k = 0; k < m; ++k) {
                sum += A[i * m + k] * B[k * p + j];
            }
            C[i * p + j] = sum;
        }
    }
}

Code Gen

With our IR in place, we can now do the final step and convert into CUDA. In the make_kernel, you can see 9 places where we generate code based off what the graph looks like (look for the format! macro in the code). We’ll go into only the rules we need for our naive matmul example for now. Feel free to join our Discord if you’d like to ask more questions.

Nodes in Graph

We start off our code gen pass with a topologically sorted compute graph. Below shows our simple representation for a 2x2 matmul.

NODE: Tensor { name: "A" } Loop Level: 0
NODE: Tensor { name: "acc" } Loop Level: 0
NODE: Tensor { name: "B" } Loop Level: 0
NODE: LoopIn { range: "2", stride: "z * 2" } Loop Level: 0
NODE: LoopIn { range: "2", stride: "z" } Loop Level: 1
NODE: LoopIn { range: "1", stride: "0" } Loop Level: 2
NODE: LoopIn { range: "1", stride: "0" } Loop Level: 3
NODE: LoopIn { range: "1", stride: "0" } Loop Level: 4
NODE: LoopIn { range: "1", stride: "0" } Loop Level: 5
NODE: LoopIn { range: "2", stride: "Acca" } Loop Level: 6
NODE: Mul Loop Level: 7
NODE: Add Loop Level: 7

Loop level is how we keep track of what code is contained within certain loops. We start with loop level 0 when we load in the data and then progress up to loop level 7 where all of the mathematics operations happen in the most nested loop.

With this graph setup, we then generate the code based off what the graph term is and the loop level.

Tensors

Our first three nodes are Tensors (A, acc, B). These are kept in a separate array to be added to the parameter list of our kernel. No kernel lines are needed yet — they’re simply mapped to input variables with pointer flags.

LoopIn

Our first LoopIn is at loop level 0. When the loop level is smaller than the number of dimensions CUDA lets us parallelize over (6 dims: 3 for grid and 3 for threadblock), we take advantage of GPU’s built-in parallelization. For level 0 with stride “z * 2”, we generate:

int loop_f = blockIdx.x;
float* g = a + loop_f * 2;
float* h = d + loop_f * 2;

The next LoopIn is at level 1 with stride “z”, so we assign it to blockIdx.y and create another strided pointer:

int loop_i = blockIdx.y;
float* j = b + loop_i;
float* k = h + loop_i;

ZeroStrideLoop

The next four LoopIn nodes have stride “0”, creating a special case for us. When you have a stride of 0, you are never actually changing the data you use. As a consequence, our system simply reuses the existing input pointers and thus achieves high performance by doing less.

Note, there are parts of the code base that ascribe more functionality to ZeroStrideLoop, but for our naive matmul this is all that will happen.

LoopIn Acca

This loop is at level 6, meaning it runs at the thread level. Here we generate the full for loop syntax:

for (int loop_l = 0; loop_l < 2; loop_l += 1) {
                float* m = g + loop_l;
                float* n = j + loop_l * 2;

Because our stride contains “Acca”, we also create an accumulator variable:

    float e = initial_value;  // accumulator initialization

Mul and Add

Our final two nodes are at loop level 7, meaning they run inside all the nested loops. The Mul node generates:

    float p = *n * *m;

The Add node performs the accumulation:

    float q = o + p;  // accumulate the multiplication result

After processing the loop body, the accumulator result is saved back and the loop is closed:

    *k = o;  // save final result
}

Final Generated Code Structure

The complete generated kernel follows this pattern:

// Grid: ["2", "2", "1"] Threadblock: ["1", "1", "1"]
extern "C" __global__ void kernel0(float* a, float* b, float* c, float* d) {
        int loop_f = blockIdx.x;
        float* g = a + loop_f * 2;
        float* h = d + loop_f * 2;
        int loop_i = blockIdx.y;
        float* j = b + loop_i;
        float* k = h + loop_i;
        float o = *c;
        for (int loop_l = 0; loop_l < 2; loop_l += 1) {
                float* m = g + loop_l;
                float* n = j + loop_l * 2;
                float p = *n * *m;
                float q = o + p;
                o = q;
        }
        *k = o;
}

Closing

After reading this blog, you should have a clear idea on how Luminal represents naive matmul internally and then how we generate the code from this IR.

Hopefully this gives you a better idea of how we use our IR to discover non-trivial optimizations like Flash Attention.

Its an exciting time to be building!