← All notes

The bottleneck wasn't the weights: activation traffic in a WGSL matmul

When you sit down to optimize an engine running a large model in the browser, intuition tells you where to start: the weights are enormous, so the problem must be reading the weights. That's where the effort goes — compress them harder, dequantize them better, touch them as little as possible.

We measured the actual breakdown of memory traffic per generated token, and it came out the opposite of what we expected. Reading weights was the small part. The expensive part was pulling in the activations, the intermediate values the model passes from layer to layer. For every byte of weight, the engine moved tens of bytes of activation.

That redraws the whole map. Almost everything intuition was proposing attacked the small side.

The optimization that was already shipped and showed nothing

There was one improvement that did attack the right side: reading those activations four at a time instead of one at a time. It was written, it was deployed, and it worked.

Except it showed up during decode and did nothing at all during prefill. And prefill is exactly the part the user experiences as "it's frozen": all the time that passes before the first character appears.

The explanation wasn't in the optimization's code, it was far away from it. For prefill, the engine grouped several tokens per workgroup — a workgroup being the bundle of threads the GPU runs together. That grouping existed for a sensible reason: reuse the weights across several tokens at once.

But the grouping and the four-at-a-time read shared the same internal switch. Turning one on turned the other off. And it turned off two more fast paths along the way.

So: we were saving on the small side of the traffic in exchange for giving up all three optimizations on the big side. A bad trade, and it survived unnoticed because each piece made sense on its own.

Dropping the grouping left the cost of prefill at almost half. We measured it in two clean, independent sessions, and we also checked that the output numbers still matched the reference implementation, which is the only way to know that "faster" doesn't mean "wrong".

And then it stopped compiling

With the grouping gone, one of the engine's paths stopped compiling. The cause was dumb: a line that was emitted unconditionally when it only makes sense in one of the modes.

The interesting part isn't the bug, it's what uncovered it. What uncovered it was a benchmark that loaded the model with a configuration the product doesn't use. It was green. It was certifying a path other than the one users run.

A benchmark that validates a configuration nobody runs isn't a safety net. It's decoration that gives false confidence.

The win that wasn't a kernel win

The other big gain has nothing to do with the GPU. It's structural, and slightly embarrassing once you see it.

Every conversation turn reprocessed the entire conversation from the beginning. The fifth turn paid all over again for everything the previous four had already computed. With two or three messages you don't notice; over a long chat, the product becomes unusable. And it doesn't become unusable all at once; it creeps, which is the hardest way to notice anything.

Now the internal state is snapshotted at the end of each turn and reused, so each turn processes only what's new. The effect isn't "it's faster". The effect is that the wait stops growing with the conversation.

The snapshot has a catch, and it's worth telling. The recurrent part of the model can't be rewound: its state is the result of having walked the whole sequence, and there's no way to undo the last few steps. You have to copy it as is. The good news is that it takes the same space for a short conversation as for an endless one, so copying it never gets out of hand.

The attention memory, on the other hand, doesn't need copying. Moving a counter is enough: anything past it gets discarded by the model's own mask. Two pieces of the same state, two opposite strategies.

The one we had to retract

And now the ugly part. We announced an improvement that later didn't reproduce, and we had to pull it.

The optimization didn't fail. The benchmark failed. Its noise meter compared the control arm against itself. Comparing something to itself always gives zero difference, so the benchmark declared there was no noise — and with the bar at zero, any ordinary fluctuation passed for a finding.

In case there was any doubt left, on another occasion the same experiment gave three very different results depending on how much GPU another program was using at the same time. Not one line of the engine had changed between measurements.

Since then there are two new rules. Before measuring, look at who else is using the GPU. And a benchmark that never detects noise isn't a precise benchmark: it's a broken one.

What we took away

Three things, and none of them is a kernel trick.

Measure the breakdown before choosing what to optimize. Intuition points at the biggest component, and the cost usually sits in the one that gets moved the most times.

Distrust switches that govern several things at once. That's where one good optimization turns off three others without telling anybody.

And treat the benchmark as suspect code. The two stories that went worst here, the invisible optimization and the retracted win, weren't engine bugs. They were bugs in what we were using to look at the engine.

Part of a series on building a browser-native inference engine.