Layouts as Algebra: CUTLASS & CuTe

The idea in one sentence

Every tiling trick in this part is one object called a layout, a function from coordinates to addresses written shape:stride, and CUTLASS and CuTe turn kernel writing into algebra on those functions.

The formula gets a name

Lecture 23 left you with one line. An address is a base plus each index times its stride:

Put the base pointer away. It is the same for every entry, so it carries no information about the arrangement. What is left is a function from a pair of indices to an offset, and that function is completely described by four numbers: how far each index may run, and how far one step of it costs. Write the two limits and the two costs as two tuples, separated by a colon:

That is a layout. Say it out loud as “shape four by three, stride three by one”. The shape says the first index runs over 4 values and the second over 3. The stride says one step of the first index moves 3 places and one step of the second moves 1. So the entry at $(2, 1)$ sits at offset $2 \cdot 3 + 1 \cdot 1 = 7$, and you never needed to know it was a matrix.

There is a second thing in that drawing worth slowing down for. A layout will accept a single number as well as a pair. Hand it 6 and it hands back 7. It gets there by unravelling 6 against the shape, first mode fastest: $6 = 2 + 4 \cdot 1$, so $x = 6$ means $(i, j) = (2, 1)$. This is the CuTe convention, and the documentation calls the order colexicographic, which is the ordinary column-major habit generalised to any depth of nesting.

So the shape has two jobs. It bounds the coordinate, and it translates between a flat index and a structured one. The stride has one job, and it is the only one that touches memory.

One notation, and no more transposes

Here is what that buys immediately. Take the same sixteen numbers in memory and three different promises about where they went.

Row-major is $(4, 4) : (4, 1)$. Column-major is $(4, 4) : (1, 4)$. The tiled one is $((2,2),(2,2)) : ((1,2),(4,8))$, and it is what a packed panel looks like: each two by two block is written down whole before the next one begins. Three arrangements that a library would describe with a flag, a transpose and a special code path, and they differ only in the second tuple.

You can build the third one in NumPy and check it against the drawing.

import numpy as np

a = np.arange(16)                                    # one line of 16 numbers
tile = np.lib.stride_tricks.as_strided(              # shape ((2,2),(2,2))
    a, shape=(2, 2, 2, 2),                           # stride ((1,2),(4,8))
    strides=np.array([1, 2, 4, 8]) * a.itemsize)

print(tile.strides)                                  # (8, 16, 32, 64)
print(np.shares_memory(a, tile))                     # True
print(tile.transpose(2, 0, 3, 1).reshape(4, 4))
# [[ 0  2  8 10]
#  [ 1  3  9 11]
#  [ 4  6 12 14]
#  [ 5  7 13 15]]

Nothing was copied to make the view: shares_memory says so, and the four strides are the four numbers of the layout multiplied by the size of an element. The transpose and reshape at the end only exist to print the tiled reading as a square, which is the one step that does copy.

This is also why CuTe stops saying transposed. BLAS asks whether $A$ is N or T, which tells you about a convention rather than about memory. CuTe asks which mode has stride 1 and names the matrix after it: $A$ is M-major if the $M$ mode is the contiguous one, K-major if the $K$ mode is. The CUTLASS documentation says plainly that the older flags leave you unsure which mode has the stride of one, which is the only thing the hardware cares about. Name the stride-1 mode and the ambiguity is gone.

A shape is a tree

The tiled layout above had tuples inside tuples, and that was not decoration. CuTe allows the shape and the stride to be nested to any depth, with one rule: they must have the same tuple shape, so every extent has exactly one stride sitting opposite it.

Nesting buys you the ability to say “these two indices belong together” without giving up the flat reading. Look at what happens when the nesting is regular.

The layout $((2, 3), 4) : ((1, 2), 6)$ has three modes: two elements one apart, three of those two apart, four of those six apart. Now check the arithmetic. A pair one apart occupies 2 places, and the next mode strides by exactly 2. Three of those pairs occupy 6 places, and the outermost mode strides by exactly 6. Every boundary lands where the thing below it stopped, so no boundary marks anything. The whole layout lists the numbers 0 through 23 in order. It is $24 : 1$ wearing a costume.

CuTe calls the simplification coalesce, and its rule for two neighbouring modes is one line: $s_0 : d_0$ followed by $s_1 : s_0 d_0$ becomes $s_0 s_1 : d_0$. Modes of extent 1 drop out. Anything else stays as it is. Coalescing is how you find out whether a walk is one straight run through memory, which was the whole subject of Lecture 23, and the reason a row walk beat a column walk there by a factor of about thirty.

Composition, worked by hand

Now the operation the rest of the subject is built on. If a layout is a function, two layouts compose, and the composite is a layout again. Write $A \circ B$ for the function that sends $x$ to $A(B(x))$: $B$ turns a coordinate into a position, and $A$ turns that position into an address.

The CUTLASS documentation carries one small example that is worth doing digit by digit, because the answer looks strange until you see where it comes from. Take

Both have size 12. Evaluating $B$ at $x = 0 \ldots 11$ gives $0, 3, 6, 9, 1, 4, 7, 10, 2, 5, 8, 11$, and evaluating $A$ at $0 \ldots 11$ gives $0, 8, 16, 24, 32, 40, 2, 10, 18, 26, 34, 42$. Feed the first list into the second and you get $0, 24, 2, 26, 8, 32, 10, 34, 16, 40, 18, 42$.

Twelve numbers, and the layout that produces them is

Why does the first mode split into a pair? Follow the orange cells. $B$’s first mode takes four steps of 3 through $A$‘s domain, landing on positions 0, 3, 6 and 9. $A$‘s first mode is only 6 long, so the third step falls off the end of it and into the second mode. Two steps fit; then it wraps. A single mode of extent 4 cannot describe a walk that changes stride halfway, so it becomes two: extent 2 with stride 24, because three steps of $A$‘s stride 8 is 24, and extent 2 with stride 2, because one step into $A$‘s second mode costs 2.

$B$‘s second mode has an easier time. Three steps of 1 stay inside $A$‘s first mode, so it stays a single mode and takes $A$‘s stride: $3 : 8$. Put the two together and you have the answer.

The nested shape is a record of where the walk crossed a boundary. That is all it ever is.

Twenty lines of Python are enough to check any of this, and checking it yourself is the fastest way to stop being suspicious of the notation.

def flatten(t):
    out = []
    for x in t:
        out.extend(flatten(x)) if isinstance(x, tuple) else out.append(x)
    return tuple(out)

def unravel(shape, x):
    """1-D coordinate to natural coordinate, first mode fastest."""
    c = []
    for s in flatten(shape):
        c.append(x % s)
        x //= s
    return c

def layout(shape, stride):
    """A layout is a function from a coordinate to an offset."""
    return lambda x: sum(i * d for i, d in zip(unravel(shape, x), flatten(stride)))

fA = layout((6, 2), (8, 2))
fB = layout((4, 3), (3, 1))
fR = layout(((2, 2), 3), ((24, 2), 8))

print([fB(x) for x in range(12)])
# [0, 3, 6, 9, 1, 4, 7, 10, 2, 5, 8, 11]
print([fA(fB(x)) for x in range(12)])
# [0, 24, 2, 26, 8, 32, 10, 34, 16, 40, 18, 42]
print([fA(fB(x)) for x in range(12)] == [fR(x) for x in range(12)])
# True

Composition is not always defined. If $B$ asks for a walk that $A$ cannot express with strides, there is no layout to hand back, and the conditions under which it does exist are written out in the CUTLASS layout algebra reference. In practice you meet composition through the two operations built on top of it, and those two are the ones a kernel actually calls.

Divide is tiling

Ask the question a kernel actually asks. Here is a matrix; here is a tile size; give me the piece of the matrix that belongs to this block. In layout language that is a division.

logical_divide takes a layout and a tiler, and returns a layout with two modes: the first says where you are inside a tile, the second says which tile you are in. It is defined as a composition, $A \circ (B, B^{*})$, where $B$ is the tile and $B^{*}$ is its complement, the layout of everything $B$ misses. The complement is what makes the second mode come out right, and it is the reason you do not have to compute any of the index arithmetic yourself.

Read the result off the drawing. Inside a tile you move 1 down a column and 6 across, because that is what the parent layout does. Between tiles you move 3 down and 12 across, because a tile is 3 tall and 2 wide and the parent’s column stride is 6. Every one of those numbers is a product of two numbers you already had. Nothing about the tile size is hard-coded anywhere.

The mirror operation is logical_product, which goes the other way: give it a tile and a pattern of repetition and it lays copies of the tile out according to the pattern. Divide cuts a big thing into small ones. Product builds a big thing out of copies of a small one. A kernel needs both, because the data arrives as a big thing and the instruction it will run on is a small one.

The threads are a layout too

Here is the move that makes the algebra pay. A thread index is a coordinate. Nothing else. So the assignment of work to threads is a layout, and it composes with the layouts already in play.

CuTe writes these as thread-value layouts. The domain is a pair, which thread and which of that thread’s values, and the codomain is an offset inside the tile. Two of them, for the same 32 numbers and the same 8 threads:

The left layout is $(8, 4) : (1, 8)$: thread $t$ takes offset $t$, then $t + 8$, then $t + 16$. On any single step the eight threads want eight offsets in a row. The right layout is $(8, 4) : (4, 1)$: thread $t$ takes a run of four to itself, so on any single step the eight threads are four apart.

That is coalescing, which Lecture 25 introduced as a rule about warps and cache lines, restated as one number in a stride tuple. The hardware fact has not changed. What changed is that you can now write down the good pattern and the bad one in the same notation, hand either to the same copy routine, and compare them without rewriting a loop.

The kernel, in three lines

Put the pieces together and the hierarchy that Lecture 26 built by hand, block tile then warp tile then thread tile, is a divide followed by a partition.

CUTLASS ships this as a tutorial kernel, examples/cute/tutorial/sgemm_1.cu. Here is the part that does the partitioning, with the shared memory declarations and the comment banners taken out:

// Get the appropriate blocks for this thread block
auto cta_coord = make_coord(blockIdx.x, blockIdx.y, _);              // (m,n,k)
Tensor gA = local_tile(mA, cta_tiler, cta_coord, Step<_1, X,_1>{});  // (BLK_M,BLK_K,k)
Tensor gB = local_tile(mB, cta_tiler, cta_coord, Step< X,_1,_1>{});  // (BLK_N,BLK_K,k)
Tensor gC = local_tile(mC, cta_tiler, cta_coord, Step<_1,_1, X>{});  // (BLK_M,BLK_N)

Tensor sA = make_tensor(make_smem_ptr(smemA), sA_layout);            // (BLK_M,BLK_K)
Tensor sB = make_tensor(make_smem_ptr(smemB), sB_layout);            // (BLK_N,BLK_K)

Tensor tAgA = local_partition(gA, tA, threadIdx.x);                  // (THR_M,THR_K,k)
Tensor tAsA = local_partition(sA, tA, threadIdx.x);                  // (THR_M,THR_K)
Tensor tBgB = local_partition(gB, tB, threadIdx.x);                  // (THR_N,THR_K,k)
Tensor tBsB = local_partition(sB, tB, threadIdx.x);                  // (THR_N,THR_K)

// Partition sA (BLK_M, BLK_K) by the rows of tC
Tensor tCsA = local_partition(sA, tC, threadIdx.x, Step<_1, X>{});   // (THR_M,BLK_K)
// Partition sB (BLK_N, BLK_K) by the cols of tC
Tensor tCsB = local_partition(sB, tC, threadIdx.x, Step< X,_1>{});   // (THR_N,BLK_K)
// Partition gC (M,N) by the tile of tC
Tensor tCgC = local_partition(gC, tC, threadIdx.x, Step<_1,_1>{});   // (THR_M,THR_N)

No GPU ran that here, and none of the numbers in this chapter came from one. Read it as a teaching artifact. The shape of the code is the point, and there is no benchmark behind it.

local_tile applies the tiler and slices into the second mode, the one that says which tile. It hands back this block’s piece of $A$, $B$ and $C$. The Step argument says which modes of the tiler to use, so the same call cuts $A$ along $M$ and $K$ while cutting $C$ along $M$ and $N$. Then local_partition slices into the first mode instead, using the thread layout, and hands back this thread’s piece. The comments give the resulting shapes, and the trailing k on gA is the count of K-tiles still to be looped over.

The inner loop is then two copies and a multiply, on tensors whose shapes already say who owns what:

for (int k_tile = 0; k_tile < K_TILE_MAX; ++k_tile)
{
  copy(tAgA(_,_,k_tile), tAsA);      // A   (THR_M,THR_K) -> (THR_M,THR_K)
  copy(tBgB(_,_,k_tile), tBsB);      // B   (THR_N,THR_K) -> (THR_N,THR_K)

  cp_async_fence();
  cp_async_wait<0>();
  __syncthreads();

  gemm(tCsA, tCsB, tCrC);            // (THR_M,THR_N) += (THR_M,BLK_K) * (THR_N,BLK_K)

  __syncthreads();
}

Compare that to the ten rungs of Lecture 26. The rungs have not gone away. Coalescing is the stride in the thread layout. Shared memory blocking is the tile size handed to local_tile, and the results-per-thread decision is the shape of tC. What has gone away is the index arithmetic, and with it the class of bug where a kernel is correct at one tile size and silently wrong at another.

The landscape, honestly

CUTLASS is a C++ template library for GEMM and the operations around it. CuTe is the layout core it was rebuilt on, and it is the part worth learning even if you never write a CUTLASS kernel, because the ideas are portable and the templates are not. There is a Python interface now, the CuTe DSL, which expresses the same algebra without the C++.

Triton takes a different bet. You write a block-level program in Python, describe what a block of the output should be, and the compiler chooses the thread mapping. You give up the ability to say exactly which thread holds which value, and in exchange you stop having to say it. NVIDIA’s CUDA Tile, whose Python form is cuTile, sits in the same neighbourhood: its own documentation says tile programs “express block-level parallelism only with no exposure to individual threads within the block”, and that a tile is a unit of data while a block stays the unit of execution.

The reason all of these exist is a hardware trend, and the trend has one direction.

Volta gave the matrix multiply its own instruction, executed by a whole warp. Ampere added cp.async, so a copy from global to shared memory no longer passes through registers. Hopper added the Tensor Memory Accelerator and warpgroup MMA: Colfax Research’s TMA tutorial states that “for a TMA copy, only one thread will be responsible for issuing the TMA operation”, and their WGMMA tutorial gives the warpgroup as 128 contiguous threads with operand $B$ always in shared memory. Lecture 26.5 takes those two apart and builds the pipeline they make possible.

Blackwell continues it. In Colfax Research’s tutorial on writing GEMM kernels with Tensor Memory, Tensor Memory is 256KB per SM organised as 512 columns by 128 lanes of 32-bit cells, the tcgen05.mma instruction they call UMMA is launched by a single thread, and the accumulator must live in Tensor Memory rather than in registers. Their summary of the arc is the sentence to keep: Tensor Memory and UMMA “do for MMA just what TMA did for copy, making it a single-threaded, asynchronous operation that does not consume registers”.

Each generation makes the thread a worse unit to think in and the tile a better one. A programming model built on layouts was going to win that argument eventually, because a layout is exactly the thing that survives when the thread stops being where the data lives.

Where this is going

This part started with an address and ends with an algebra of addresses. Lecture 23 said a matrix is a promise about where its entries went. Lecture 24 turned seconds into milliseconds by changing only the order that promise is walked in. Lectures 25 and 26 said the same thing again with thousands of threads. This chapter says all of it is one type of function, closed under the operations you need.

The course opened with a smaller question: is $b$ in the span of the columns of $A$? The move that answered it was to name a structure and then compute with it. Layouts are that move applied one level down, to the addresses those columns live at. The algebra of $Ax = b$ tells you what the machine must compute. The algebra of layouts tells it where to stand.

One chapter is left in this part, and it drops an assumption everything so far has leaned on. Every layout here is a promise you can evaluate: give it a coordinate, get an address, no lookup required. When almost every entry of a matrix is zero, that promise stops being worth keeping, and the address has to be stored instead of computed. Lecture 28 is about what that costs.

If you want to write the code, the Library has the three places worth starting: Triton Puzzles for block-level thinking, LeetGPU for CUDA practice judged on real hardware, and the GPU MODE kernel leaderboard for problems that people are still competing on.