Stumbling Through the GPU

Stumbling Through the GPU

I have a story about something that is, as it stands, almost entirely useless. I've harnessed the computing power equivalent of my birth-years top supercomputer to have a model decide where to move a visually distinct and non-infringing polyomino ^{1certainly bears no resemblance to the well defended intellectual property of Tetris Holding LLC, which I will refer to as REDACTED}. While the task explored in the article that follows does not derive significant shareholder value^{2my wife was distinctly non-plussed at seeing what I'd spent so much of my time on}, I would say that it will be very useful for features to come in the near future. So, if you have nothing better to do, please join me for a journey through packign weights, shared memory, synchronization barriers, matrix instructions, and a particularly persuasive result that mostly disappeared when we changed the prompt length.

Getting a Baseline

The experiments used our ^{3do I get to say 'our' because I made a fancy website name with "labs" in it} native wgpu runtime on an M1 Max, running an NF4-quantized ^{4A floating point format with specially placed values that give decent precision for much smaller size} 250M classifier^{5I'll be releasing some more information about the model creation later down the line}. The model is based on one from liquid labs with some significant changes to its structure that, along with the quantization, bring the whole model into the ~60MB range without collapsing its ability to play this little game.

Spweing its first decision took 320 milliseconds, even GPUs need to warm up, and the next 2 decisions took 269 and 219 milliseconds. Somewhat underwhelming with human [REDACTED] players averaging around 2-3, but running on a users device thats what I'd call "serviceable" with room for improvement.

With our baseline set at about ~2-3 Decisions per second, ~180 Actions Per Minute (APM^{6Not the [REDACTED] native Attacks Per Minute term}), we can get off to the optimization races.

One request still has plenty of work inside it

The prompt always has 346 tokens in it, and looks like this:

<|startoftext|><|im_start|>system
You are playing a polyomino game. The board state is below.
. is an empty cell
# is a settled cell
% is your falling piece<|im_end|>
<|im_start|>user
<board>
. . . . . . . . . . .
. . . . . . . . . . .
. . . . . . . . . . .
. . . . . . . . . . .
. . . . . . . . . . .
. . . . . . . . . . .
. . . . . . . . . . .
. . . . . . . . . . .
. . . . . . . . . . .
. . . . . . . . . . .
. . . . . . . . . . .
. . . . % % . . . . .
. . . . % . . . . . .
. . . . % . . . . . .
. . . . . . . . . . .
. . . . . . . . . . .
. . . . . . . . . . .
. . . . # # . . . . .
. . . . . # # . . . .
. . . . . . # # # . .
. . . . . . . # . . .
</board>
piece:J4 rot:1 x:04 y:12
next:I5
heights:00 00 00 00 04 04 03 02 02 00 00
holes:07 drop:03<|im_end|>
<|im_start|>assistant
action:

We play the game one simulation step at a time, which means we're working at a batch size of 1. Each step, we're predicting if the AI should take one of 6 actions: Soft Drop, Hard Drop, Rotate CW, Rotate CCW, Move Left, Move Right. I've reserved 2 more outputs for future experimentation, and I like powers of 2, so the total is 8 but 2 are never used.

For the convenience of those mostly familiar with LLMs and those more familiar with traditional ML, we can think of this problem in two ways:

Classification: This is a classification problem with 6 potential output classes

Language Modeling: We are autoregressively predicting 1 next token in a sequence after doing ~350 tokens of prefill. Prefix caching is far less useful because of how the prompt is structured, so we're focused on prefill speed

Processing a whole prompt and processing 1 new token give the GPU very different jobs. A kernel that spreads a large matrix across the machine can leave a small matrix paying for arrangements it barely uses.

Slow and Steady

For testing each iteration I had 2 helpers: a kernel harness to time individual GPU ops, and the classifier to measure complete decisions. To make sure the math wasn't getting mangled, I also compared output scores against 3 reference cases from my full precision parent. The expected decisions were LEFT, LEFT, and CW, which turned out to be somewhat unfortunate because of some built-in biases.

Pulling the Thread

A feed-forward block takes each token's hidden dimension (1024 for LFM 2.5), stretches them out to a new hidden dimension (LFM2.5 uses 2560) with a matmul, runs them through a nonlinear function, and squashes them back down to the original size again. The stretch is 2 projections run side by side (the paired projection), which produce 2 sets of 2560 numbers^{7called gate and up, because naming things is hard}. SwiGLU, a combination of the activation functions "Swish" and Gated Linear Unit (GLU), is the combination of those operations.

gate   = x * gate_weights       # stretch: 1024 → 2560
up     = x * up_weights         # stretch: 1024 → 2560
hidden = silu(gate) * up        # combine: 2560
output = hidden * down_weights  # squash:  2560 → 1024

The trouble was in the squash; the GPU splits the down projection's 1024 outputs into 32 groups of 32 to work on the groups in parallel and, because the original kernel never stored hidden, each group of 32 read the original gate and up projections to calculate silu(gate) * up for itself. This meant the same 2560 numbers got calculated 32 times!

You can imagine the GPU splitting it up like this:

Kernel 1:
  gate   = x * gate_weights       # stretch: 1024 → 2560
  up     = x * up_weights         # stretch: 1024 → 2560

Kernel 2:
  hidden = silu(gate) * up        # combine: 2560
  output = hidden * down_weights  # squash:  2560 → 1024

Profiling (the good kind)

Measuring the pieces: each of the 14 feed-forward blocks spent about 4.43 milliseconds on its paired projections and 3.72 on its down projection. Across all 14 blocks that's roughly 114 milliseconds, with attention adding another 28ms.

My first attempt was to calculate hidden once, park it in a GPU buffer, and let all 32 groups read it from there. This cost an extra dispatch, a launch of a 'kernel' or tiny program over a grid of work, but got rid of the other 31 copies. This seems like it should work, but for the GPU it's one more thing to schedule even if the math inside is tiny and theres a new buffer R/W operation which transits memory.

The new split from the GPU's point of view was something along the lines of:

Kernel 1:
  gate   = x * gate_weights       # stretch: 1024 → 2560
  up     = x * up_weights         # stretch: 1024 → 2560

Kernel 2:
  hidden = silu(gate) * up        # combine: 2560 and store it for K3

Kernel 3:
  output = hidden * down_weights  # squash:  2560 → 1024

Surprise surprise, it did basically nothing. A newly compiled run took 3.94 milliseconds against the original's 3.72, another took 3.77 against 3.73. The outputs were bitwise identical, so I had successfully made the same answer arrive at approximately the same time. Very nice, I like.

Around this time I was browsing around and came across the concept of a fused kernel, which basically means trying to minimize the amount of time data has to be moving, or conversely, maximize the number of operations you can do on data while you have it in the super-speedy registers already.

So, attempt two was pushing the SwiGLU ops into the end of the paired projection kernel where gate and up were already sitting around waiting to be consumed. Now instead of data having to move in and out for each kernel, we do a single combined buffer write and the intermediate storage drops from 7.09MB to 3.54MB with no extra dispatch. Roughly:

Kernel 1:
  gate   = x * gate_weights       # stretch: 1024 → 2560
  up     = x * up_weights         # stretch: 1024 → 2560
  hidden = silu(gate) * up        # combine: 2560

Kernel 2:
  output = hidden * down_weights  # squash:  2560 → 1024

The 2 kernels together went from about 8.34 to 8.30 milliseconds. Yes, I'm keeping the second decimal place because without it the improvement disappears entirely and I want to feel good about myself. Still, the smaller buffer was nice and the next kernel got a simpler input, so this version stayed.

Naturally I decided memory traffic was still problem anyway because I'd read it on someones blog and it sounded nice. What I'd actually proven though was that shuffling this one activation around didn't change much. The matrix multiplication, shared-memory access, and scheduling were all still in there somewhere declining to admit themselves as the problem.

Unpacking the Weights

A GPU workgroup is a bundle of little worker bees with a shared memory pool. Each workgroup is responsible for some section of the whole models compute work: it loads some of the models activations, decodes the quantized weights into FP32 or BF16 and uses those values to calculate its share of the output. When the workgroup completes a chunk it grabs the next one.

For our current task, imagine the output as a matrix with a row for each token and a column for each output feature. My shader chops that into tiles of 32 rows by 32 columns and hands each tile to a workgroup.

There had to be some duplicated work somewhere, and my first theory was that each token row was unpacking the same weights over and over again. Reading the shader I found that I was, unfortunately, already sharing the decoded weights across all 32 rows within a tile.

Digging through it some more, though, I found there was some repetition happening across workgroups; a 346-token prompt needs 11 tiles of 32 rows, so every NF4 quantized weight gets decoded into FP32 11 times by groups working on different token positions. The activations get the same treatment along the other axis: 2560 output columns in 32-column tiles means 80 groups reading the same input values.

Some of those reads will come from cache. That's nice, though we're still asking for the same things repeatedly and I would prefer that we didn't. So, we can either share more data within each workgroup, or unpack all the weights once and let the groups read them already unpacked.

For the second option I already had a dense float32 matrix multiplication kernel sitting around, so I gave it a go. Across the shapes I tested it was 4.5 to 4.9 times slower than packed NF4, and keeping all the feed-forward weights unpacked would mean another 367.5MB of floats hanging around in memory. Somewhat inconvenient for a model whose ~60MB size I'd been feeling good about^{8I was hoping to make the little model faster without having to explain why it had become a much bigger model.}.

There is a catch in this comparison: a dense kernel and a packed kernel can divide up the work differently. Our kernels use different tiles and load data differently, so this comparison only really told us that my existing dense kernel was a bad destination for the unpacked weights. Giving compression all the credit for that difference would've required a comparison where the rest of the work was the same.

Our ternary path was also significantly slower in several tests, despite having simpler values to decode and my very reasonable expectation that smaller weights should make things faster^{9The GPU continues to have very little interest in how nice I think the storage format is.}.

Bigger Tiles

Each tile works through a matrix' inner dimension^{10The dimension we sum over when we multiply a row by a column.} in chunks, and K_STEP sets the size of a chunk; between chunks a workgroup has to synchronize, which means waiting for shared data to be ready before reading it, and for everyone to finish using it before overwriting it with the next set.

Naturally, waiting seemed like something we should do less of, so I naively bumped K_STEP from 32 to 64 to get twice as much work done between synchronization rounds. A chunk we're keeping around is now twice as big, so this needs more shared memory, but cutting the number of rounds in half seemed worth it.

The paired projection went from 4.62 milliseconds to 13.35, almost 3 times slower. The down projection went from 3.72 to 5.90. I had successfully reduced the amount of waiting, at least in the shader, and considerably increased the amount of waiting for me, yipee.

There were several possible explanations: a bigger allocation might leave room for fewer workgroups to run at once, or register pressure and compiler changes could be responsible. Without counters to tell which of those was happening, reverting this change was easy enough, so that's what we did.

So, attempt two: making the tile taller instead, giving a workgroup 64 token rows to work on and letting it use the same decoded weights across more of the prompt. The paired projection dropped to about 3.88 milliseconds and the down projection to 3.02. Very nice, an improvement that survives rounding.

Unfortunately the shorter prompts weren't having nearly as much fun; at 64 rows one projection went from 0.82 milliseconds to 1.75. A group is now working through more rows on its own. With a small input there aren't enough groups to keep the rest of the GPU busy.

I already had a kernel for small row counts, so after comparing the two I set a cutoff at 96 rows: anything below that uses the short-row path, and anything from 96 up gets the bigger tiles. I also checked the outputs at 95, 96, and 97. I know who wrote the condition and would prefer to catch his mistakes here.

Then I tried going the other way with K_STEP, shrinking it to 16. Compared with the taller tile's K32 version, this halved the staging memory^{11The shared memory holding a chunk of weights and activations while the group works on it.}, doubled the number of synchronization rounds, and made the paired kernels another 8% to 10% faster.

So far, synchronizing less had made things considerably worse and synchronizing more had made them better. My initial theory about waiting was having some difficulties. The amount of data we were holding onto seemed relevant, though I'd been quite confident about the other idea too.

A Very Convincing Wrong Turn

Next I tried a wider tile, 32 rows by 64 columns. This lets each group share its activations across more output columns, but means more row groups unpacking the weights again, and at our usual 346 tokens it looked promising: one projection went from 2.89 to 2.78 milliseconds, and another went from 2.87 to 2.47.

I liked this because I had an explanation involving activation reuse and compressed weight sizes, and the numbers were being very accommodating. Then I changed the prompt length.

The trouble is that we have to schedule whole tiles, so a prompt with 346 rows needs 6 tiles of 64 rows, with room for 384, or 11 tiles of 32 rows, with room for 352. My shader checks the edges to keep those extra positions out of an answer, but the GPU still gets handed those differently shaped piles of work.

At the particular length I'd been testing, the wide version needed 176 workgroups and the narrow one needed 192, giving it about 8% less padding before any of my clever ideas about memory reuse had a chance to contribute. When I tried some more prompt lengths the narrow tile won or tied at every one I tested, including 346. I deleted the wide version and went back to 64-by-32 with K16.

I'd spent a fair amount of time coming up with an explanation for what appears to have been padding and measurement noise, which is one of the advantages of always testing the same little game: you can get very good at the dimensions of that particular board description^{12Optimizing for the rounding involved in dividing 346 by 64 was not quite the GPU expertise I had hoped to acquire.}.

Surely the Matrix Instructions Will Help

Around this time I found that the M1 Max has native cooperative matrix operations, which let subgroups work together on small matrix fragments. Since most of what I was doing was multiplying matrices, I thought instructions specifically for that purpose should be useful and put together a prototype using 8-by-8 float fragments with the same K16 staging loop.

Getting a correct set of numbers out of this took some work because my staged tiles were row-major and one of the load operations expected column-major^{13These describe whether adjacent values in memory belong to the same row or the same column. Mixing them up has a fairly predictable effect on matrix multiplication.}. The output scores were wrong, but the chosen actions were LEFT, LEFT, and CW.

All 3 had picked the expected action with the matrices loaded sideways! Remember those built-in biases from earlier? The model was sufficiently enthusiastic about going left that it could survive this and still make the same decision, which was a good reason to keep comparing the actual scores^{14Checking only the winning class would've saved me quite a lot of work here, with some consequences for the usefulness of the result.}.

The shader validator also had an objection: cooperative operations need every participating invocation to reach them together, with uniform control flow. Some of the values and helper calls looked perfectly uniform to me. To convince the validator I rearranged the staging guards and changed how the groups loaded their shared control values.

With all of that sorted out, the correct version was 19% to 28% slower across the shapes I tested. A bigger staging chunk made it worse. This float32 implementation spent whatever we'd saved with the matrix instructions somewhere in the decoding, staging, synchronization, and edge handling, so it wasn't doing us any favours.

I removed the execution path and kept the shader around in case a different precision or layout makes it useful later. For now it can sit there with the other things I was excited about before I measured them.

Smaller Temporary Values

The workgroups were still using a lot of shared memory, so I tried a float16 version of the stored weights, activations, or both while keeping decoding, accumulation, and output in float32. A float16 value is smaller, but converting into and out of it adds work and some rounding, which we would have to check along with our timings.

I alternated between the versions and saw fairly consistent gains of about 1.5% to 2.7% for plain matrix multiplication when I staged both weights and activations in float16, then checked aligned row counts and got similar gains, so at least this time the improvement survived having a prompt that fitted neatly into the tiles.

The fused operations were less consistent. One path that calculated its activation inside the consumer actually got slower. These intermediate values needed their own test tolerance: maximum absolute differences reached about 0.011 on matrix-multiply outputs and 0.27 after the SwiGLU product^{15How much difference the test allows before it decides I've broken the math.}.

A couple of percent on one path, a regression on another, and another numerical tolerance to keep track of seemed like a fairly modest reward for a change like this, so float16 staging stayed behind a development switch.

While reviewing this I also found that my occupancy math was wrong: I'd divided the adapter's workgroup-memory limit by each group's usage to work out how many groups could be resident at once. Unfortunately that limit only tells us how much 1 workgroup can allocate. It says nothing about the per-core memory pool or how many groups will actually fit in it.

The trace I had covered whole command buffers, with lots of dispatches in each submission, so I'd need a more detailed capture to get the counters for this. Smaller staging had helped, but shared-memory access, arithmetic, occupancy, or some combination of them could still explain the result. My spreadsheet had been very specific about something it wasn't measuring^{16I have since corrected the spreadsheet, which is considerably easier than getting the missing GPU counters.}.

Do We Need All of These Rows?

After spending this much time inside the kernels I went back to the executor. This had a simpler problem: the final feed-forward block was processing all 346 rows, and our classifier was only reading the last one to make its decision.

Nothing after that block needed the earlier outputs either, so I sliced the input and residual down to 1 row and sent the final block through the single-row path. Attention and convolution still need the whole prompt and their state for continuation. In this last feed-forward block we could stop calculating the other 345 rows that I'd been throwing away.

When I timed the decisions again they came down to 223, 185, and 171 milliseconds, against the original 320, 269, and 219, with the output scores moving by a few millionths. I checked that processing a prompt in bulk and then appending gave the same result as processing its tokens one at a time, so continuation was still working.

One of those drops was bigger than removing a ~7-millisecond operation could explain. Changes to dispatches and allocations might account for some of it, with run-to-run variation doing some more. I'd only timed 3 cases once each, so this gives us 3 timings and an encouraging result, with many more runs needed to describe a latency distribution^{17Tempting as it is to give 3 numbers a chart and start referring to them as a study.}.

Remember when I said prefix caching was far less useful because of the prompt structure? It turns out there are still 47 shared tokens before the board state changes, so we cached those and had each new decision branch off a shared base. A new board starts from that base every time, which means our model doesn't accumulate a conversation about everywhere the piece has been.

In that session the full decisions took 0.61 to 0.68 seconds and the cached-tail decisions took 0.53 to 0.56, about 15% to 18% faster, with the shared-base append test matching the output scores exactly.

You may have noticed those times are considerably worse than our timings a few paragraphs ago, which had me looking through the classifier harness again. It was submitting work, polling, and sleeping for 1 millisecond between polls across roughly 30 sequential operations, and with the host under load each wakeup picked up some scheduler delay while the kernel timings stayed about the same.

I can compare a cached run with an uncached run from that session. Comparing them with the earlier quiet run would mostly tell us what else my computer was doing. Apparently I had found another way to measure waiting.

Rows still have to fit into whole tiles: dropping 47 tokens leaves 299 fresh rows, which fit into 5 tiles of 64 instead of 6. That's 13.6% fewer tokens and 16.7% fewer row tiles, which is fairly close to the gain we measured even if it doesn't account for every part of the runtime.

Moving more of our stable information to the front of a prompt could give us more to cache, though changing the prompt also changes what the model sees and would need another quality check, since I'd like the next speedup to preserve its current level of expertise at moving left.

Where That Leaves the Little Blocks

The kernel now uses 64-by-32 tiles with K16 staging from 96 rows up. Below that we use the short-row path. SwiGLU stayed at the end of the paired projection, the final feed-forward block only calculates the row we read, and I've deleted the wide tile and cooperative matrix path while leaving float16 staging behind its development switch.

The later decisions in the final-row experiment took 185 and 171 milliseconds against the original 269 and 219, and prefix caching gave us another comparison with a 15% to 18% improvement in its own session. I made those comparisons in separate sessions, so adding those percentages together would require a level of enthusiasm for the results that even I can't justify.

For our [REDACTED] purposes, 185 and 171 milliseconds works out to roughly 320 to 350 APM, up from the ~180 we started with. I'm using 2 measurements here. This should probably discourage anyone from putting it on a leaderboard^{18I would be very interested in the category for a 250M classifier on an M1 Max, though.}.

There's still plenty to investigate in the GPU, since I haven't actually measured its main limiter and the explanations I've come up with are still waiting for evidence. The smaller kernel tweaks were giving me less each time, but skipping repeated work had given us something useful to build on.

I'd also like to try the ternary weights again because their codes take half the bits of NF4, though this kernel still expands them into float tiles and does all the same activation reads, multiplication, and synchronization. To see what the format itself buys us I'll need to encode the same representable weights both ways and use the same tile geometry, so the comparison has fewer other things to blame.

For now our little game makes its decisions faster, our runtime does a bit less repeated work, and I've deleted several things that seemed like very good ideas when I wrote them. This should all be useful for the product features coming next, and in the meantime we can enjoy the polyomino getting to the left slightly sooner.

We already know how much it likes going left, because I tried loading its matrices sideways and it carried on anyway.

---

These measurements used 1 classifier on an M1 Max through the native wgpu path.

Notes

1. certainly bears no resemblance to the well defended intellectual property of Tetris Holding LLC, which I will refer to as REDACTED ↩
2. my wife was distinctly non-plussed at seeing what I'd spent so much of my time on ↩
3. do I get to say 'our' because I made a fancy website name with "labs" in it ↩
4. A floating point format with specially placed values that give decent precision for much smaller size ↩
5. I'll be releasing some more information about the model creation later down the line ↩
6. Not the [REDACTED] native Attacks Per Minute term ↩
7. called gate and up, because naming things is hard ↩
8. I was hoping to make the little model faster without having to explain why it had become a much bigger model. ↩
9. The GPU continues to have very little interest in how nice I think the storage format is. ↩
10. The dimension we sum over when we multiply a row by a column. ↩
11. The shared memory holding a chunk of weights and activations while the group works on it. ↩
12. Optimizing for the rounding involved in dividing 346 by 64 was not quite the GPU expertise I had hoped to acquire. ↩
13. These describe whether adjacent values in memory belong to the same row or the same column. Mixing them up has a fairly predictable effect on matrix multiplication. ↩
14. Checking only the winning class would've saved me quite a lot of work here, with some consequences for the usefulness of the result. ↩
15. How much difference the test allows before it decides I've broken the math. ↩
16. I have since corrected the spreadsheet, which is considerably easier than getting the missing GPU counters. ↩
17. Tempting as it is to give 3 numbers a chart and start referring to them as a study. ↩
18. I would be very interested in the category for a 250M classifier on an M1 Max, though. ↩