More Stumbling, This Time in the Browser

In the last post I had persuaded a small model to play our visually distinct polyomino game a bit more quickly, and spent quite a lot of time explaining improvements that disappeared when I changed the prompt length. I'd like to stop doing the second part, though it seems to come quite naturally.

Last night I worked somewhat inefficiently through another 25 ideas and kept 4: a dense matrix kernel that lowered its classifier's average decision time by 16.3%, an encoder change that nearly halved the memory our runtime accounted for, and 2 smaller changes on the host side. I wont go into depth on the failed ones for now because they were probably what someone familiar with GPUs wouldn't even think to try, but I did leave a brief description of them at the end for any poor soul that makes it through.

The model weights stayed the same, and all 4 survivors returned exactly the same output scores. After the sideways-matrix incident, I've become quite attached to checking those.

This time, in a browser

The previous experiments ran through native wgpu on my M1 Max, but our browser runtime has a few more things between the Rust code and an answer: WebAssembly, JavaScript, WebGPU, the browser's GPU implementation, and whatever the driver makes of all that. The polyomino has acquired quite an entourage.

This is where the feature will actually run, so I timed every candidate that reached that stage with a complete model in Chrome 154, using WebGPU on the same M1 Max. A good result in the native kernel harness could get an idea onto my list, but it had to make the browser task faster to stay.

For the dense and packed-weight paths I used our existing polyomino classifiers and their 346-token game-state prompt. For the encoder I used our MagicBox ternary model, with 16 layers, a fixed 111-token input, 2 questions, and all 4 pointer projections that produce the start and end scores for extraction.

I gave each candidate its own branch, with helpers working on different implementations in batches of 5 while compilation and GPU measurements ran one at a time. This gave each result 1 source change to blame, which seemed helpful given how many explanations I'd managed to produce for the last set of results.

The loop was:

1. Compile the candidate Rust runtime to WebAssembly once.
2. Measure the original runtime on the same task 5 times for that model and batch.
3. Load the candidate runtime and run that task 5 times in the browser.
4. Compare the arithmetic means, output scores, and chosen answers.
5. Keep a correct candidate with a lower mean, or restore the original source.

I counted the first inference along with the next 4, with no warmup calls. Downloads, WebAssembly initialization, and loading the weights happened before I started the timer; inference, GPU readback, and getting the result back all happened inside it.

The classifier also tokenized its prompt, while the encoder started with fixed token IDs and included prediction and parsing the returned result, so I kept their comparisons separate. I checked hashes of the model, configuration, tokenizer where applicable, and task inputs too, since quietly swapping a checkpoint would be a fairly easy way to make a kernel appear more capable than it was.

Four answers per worker

The dense matrix kernel gave a workgroup of 256 invocations a 16-by-16 output tile, with each invocation working on 1 output value^{1An invocation is one execution of the shader. The workgroup gives 256 of them shared storage and a way to synchronize.}. Each row belongs to a token and each column to an output feature, so to calculate a cell we walk along the inner dimension, multiply an activation by its corresponding weight, and add the products together.

I gave each invocation a 2-by-2 square of outputs instead. At each step it now grabs 2 activations and 2 weights from shared storage, then uses those 4 values to update 4 accumulators:

top_left     += activation_0 * weight_0
top_right    += activation_0 * weight_1
bottom_left  += activation_1 * weight_0
bottom_right += activation_1 * weight_1

Each value gets used twice, and the same 256 invocations now cover a 32-by-32 tile, producing 4 times as many outputs per full tile and shrinking the dispatch grid.

Of course, it also has to keep 4 accumulators around, and the shared staging grows from 2 KiB to 4 KiB. I'd just spent a post discovering how much the GPU could dislike a bigger tile, so I was fairly prepared to delete this one too.

The dense classifier's 5-call mean went from 204.10 milliseconds to 170.74, a 16.34% reduction, though the first call got slower: 269.5 to 316.1 milliseconds. Calls 2 through 5 averaged 187.75 milliseconds in the original and 134.40 in my version, so this time there was a repeated-call improvement large enough for me to see without squinting at the last decimal place.

I kept the original single-row matrix-vector path and the edge masks, and each output still accumulated along the inner dimension in the same order. All 8 classifier scores matched exactly, so the bigger tile stayed.

These numbers belong to the dense floating-point model. Our packed NF4 and ternary paths use their own kernels, and as we'll get to shortly, they had rather different opinions about the changes I tried.

The encoder kept everything

An encoder layer makes quite a few intermediate matrices as it goes: normalized hidden states, projections, attention outputs, feed-forward intermediates, and residuals. My executor gave each operation a fresh output buffer and hung onto every one until the whole pass finished, which adds up rather quickly when we do it for 16 layers.

There was a reason for keeping them. The GPU runs asynchronously, so finishing the Rust function that records a dispatch tells us very little about whether the GPU has finished reading that dispatch's inputs, and releasing them then would give the model some fairly ȋ̷̼ñ̶̡͚̚t̷͓̮̓ę̴̛̬r̴̼̐ę̸͊̽͜s̷̞̳̊t̷͕͑i̸̝͛n̴̟͝g̷̡̫̔ material to work with.

The trouble being that I was also keeping values after I'd recorded their final consumer, while later layers asked for temporary storage of exactly the same shape. An old allocation was sitting there holding something we'd never read again, and I was already trying to find another place to put the next thing. I'd effectively Hotel California'd our poor matrices because I'd given every intermediate result somewhere to live but forgotten to arrange for any of them to leave.

So, I changed the task's scratch arena to keep track of when each intermediate had its final consumer recorded. At that point its allocation can go back into the pile available to later operations with the same shape, while the task keeps owning the buffer until completion.

The order matters here:

write an intermediate into buffer A
record its final consumer, which reads buffer A
record a later operation that writes a new value into buffer A

The GPU gets those commands in that order, so the earlier read happens before the later write. I kept live inputs separate from their output buffers, and the task still retains both live and reusable allocations until its fence confirms completion^{2The fence is our completion boundary. It tells the task when submitted GPU work has finished and its retained resources can be released.}.

Naturally there was an exception waiting for me: the final mask comes from a CPU upload, and a host write can reach the queue before we submit the recorded dispatches. If I reused an earlier buffer for that upload, I could overwrite an input before the older dispatch got to read it, so the mask kept fresh storage.

I also kept the existing ownership rules for cancellation, dropped tasks, failed operations, and failed fences. We can hand a buffer to a later recorded operation while the GPU still needs it; losing track of that buffer would be a considerably more adventurous change.

For the MagicBox task, peak runtime-accounted storage fell from 719.83 MiB to 374.33 MiB, a 48.0% reduction. That count comes from our allocation tracker and includes owned buffers, pooled storage, staging, uniforms, and completion-owned host results, so those are the allocations we're comparing.

The model equations and kernel dispatch counts stayed the same. We'd found somewhere to put the same calculations without keeping quite so many of their old answers around.

The 5-call mean fell from 115.54 milliseconds to 69.58, or 39.78%, which looked very nice until I started looking at where the time had gone:

Quite a lot of it had gone out of the first call. Calls 2 through 5 averaged 73.0 milliseconds against 60.5, and the final 3 averaged 64.07 against 60.40, so the timing improvement got much smaller as I worked down the table. The smaller pile of allocations, fortunately, stayed smaller whichever calls I chose.

Dividing 115.54 by 69.58 gives about 1.66, which means a 39.8% drop in time is also the equivalent of 66% more work per second in this 5-call comparison^{3Same 5 calls, now in a trench coat.}. I've kept the milliseconds beside the percentages so you can see how I got them.

A little less paperwork

Before handing work to the GPU, our backend records its operations in a list. The submission path then copied that list into a new VecDeque and made a new Vec for every uninterrupted run of compute dispatches, gathering each run into 1 compute pass.

We were making a list, copying the list, then making smaller lists. Somewhere in here there was also a GPU.

I changed it to read directly from the recorded list with a peekable iterator, stopping a pass whenever we reach a copy, clear, or mapping operation. That let me remove both temporary collections while keeping the same pipelines, buffer bindings, dynamic uniform offsets, callbacks, buffer-retirement order, and GPU dispatches.

The polyomino classifier's mean went from 146.38 to 140.48 milliseconds, a 4.03% reduction, but the first call did all the work: it dropped from 259.0 to 228.2 milliseconds. Calls 2 through 5 averaged 118.23 and 118.55, respectively. The lists were gone; the repeated calls took much the same time.

Then I put a limit on how much encoder work we record in 1 poll. The original task recorded every layer before returning control; I changed it to record up to 4 layers, submit them, wait for that fence, and continue from the residual it had kept around.

For our 16-layer model this gives us 4 layer-recording slices, with the pointer head recorded after the final layer. Waiting for each fence before continuing keeps us within the existing 1-slot completion limit and gives us each submission's errors to check.

Its 5-call mean went from 81.58 to 80.60 milliseconds, a 1.20% reduction. We're back to the decimal places, and this one's individual calls were even less accommodating: the first got slower, the final 3 averaged 63.77 milliseconds in the original against 66.13 in my version, and a much quicker second call pulled the overall mean down.

That passed the rule I'd set for the round, so I kept it as a separate candidate for review. I can see the 4-layer bound in the code, though I'd quite like a few more timings before getting attached to the 1.20%.

The attractive ideas that went away

I also timed 18 candidates that returned the right outputs and took longer to do it.

A bigger NF4 register tile took its classifier from 113.70 to 183.62 milliseconds, and specializing NF4 pipelines for our common matrix dimensions took it from 118.10 to 184.78. The specialized version's first call took 512.3 milliseconds against the original's 178.6, which gave its subsequent hundred-millisecond calls quite a lot of damage to repair in the average.

Repacking more ternary weights, caching buffer bindings, unrolling reductions, and changing the packed tile geometry all lost their browser comparisons too. I'd had explanations involving reuse, fewer lookups, or fewer loops for several of these, and I still quite like the explanations, though the answers kept arriving later.

One attention candidate didn't even get as far as browser timing. It split the key sequence into chunks, calculated local softmax statistics and value sums for each, then combined the chunks using their relative weights, which came with a rather unpleasant edge case.

A chunk's value sum could overflow before we applied its very small global weight, and with finite inputs I got a sum of infinity, a global multiplier of 0, and a merge calculating infinity * 0. The result was NaN, while the original whole-row calculation gave a finite answer for the same case.

I rejected that version. The polyomino already has enough ways to disappoint me with ordinary numbers.

Another 2 ideas turned out to be things we'd already done: creating compute pipelines lazily and keeping sampled token IDs on the GPU for the next embedding lookup. For a brief moment, reading was my fastest optimization technique.

That left 22 of the 25 ideas reaching browser timing: 18 slower candidates, 4 survivors, and 110 timed candidate tasks. The other 3 were the correctness rejection and the 2 things I could cross off by reading the code.

I kept the rejected patches, raw results, and measured packages outside runtime Git, restored their source branches, and removed the worktrees. I cleaned the shared Cargo target between builds and removed it at the end, leaving the runtime repository with the surviving source changes and their focused regression coverage.

What the numbers actually describe

I started every candidate from the same original runtime revision and applied its own change, then measured a fresh baseline for each group of experiments and model. Each row below belongs to that comparison:

The dense classifier and MagicBox encoder are doing different jobs, and even my 2 encoder changes ran separately against their own original runs. Adding the percentages together would let me describe a very impressive runtime that I haven't actually run.

All 4 candidates returned identical outputs in their browser comparisons and passed the repository's portable CPU and WebAssembly checks. You can find the source and full timings in the pull requests for dense register tiles, encoder workspace reuse, dispatch replay, and encoder recording slices.

The 2 encoder branches touch the same ownership code, so I've still got to reconcile those and run a combined build on the same browser task. Then we'll have a timing for the version we're putting together.

For now the model gets the same weights, returns the same scores, and in the dense path finishes sooner on every repeated call we measured, which should help it get on with its important work deciding where to move the polyomino.

I'll keep checking all 8 scores. Its fondness for moving left has already survived considerably worse things than a slow allocator.

Notes

1. An invocation is one execution of the shader. The workgroup gives 256 of them shared storage and a way to synchronize. ↩
2. The fence is our completion boundary. It tells the task when submitted GPU work has finished and its retained resources can be released. ↩
3. Same 5 calls, now in a trench coat. ↩