A Stolen LDS Slot Produces Wrong Tokens, Not an Error
The previous post in this series gave the headline number: mojo-baro's decode engine, running Q4_0 weights through a persistent kernel, now reads 130.7 tok/s_gen against llama.cpp's 110.0 on the same twenty prompts. This is where that number came from — five rounds run in one sitting, one of which was reverted, and one of which produced a class of bug this project had a rule against before today, and needed a new one for anyway.
The shape of a persistent kernel
kernels/mega.mojo::amar_mega_token is one grid launch that runs an entire decode step: 32 transformer layers, the final norm, the head GEMM, and the argmax, without returning to the host in between. Instead of the host sequencing thirty-odd kernel launches, ninety-six workgroups — one per compute unit on this card — synchronize against each other through a software grid barrier: every block increments a shared counter and spins, bounded, until the count says every block has arrived, then all of them proceed to the next phase together.
That design has an assumption baked into it that a normal kernel launch doesn't: every block the kernel was told exists actually gets a slot on a compute unit at the same time. A normal launch just queues blocks and lets the scheduler run them as CUs free up. A persistent kernel with a grid barrier needs all of them resident simultaneously, because a block that hasn't started yet can never increment the counter the other ninety-five are waiting on.
The question that started the day
The kernel runs at 256 VGPRs per thread by default, one block per compute unit, G=96. Dropping the register cap to 192 is a lever this project has pulled before on other kernels: fewer registers per thread can mean more threads resident, which on a bandwidth-bound kernel means more outstanding memory requests in flight. The frozen prediction for today: at the current G=96, the cap probably costs something (registers freed up don't help if there's nowhere for a second block to go); at G=192 — two blocks per compute unit, exploiting the freed registers — somewhere between a 5% loss and a 3% gain, and any run producing a wrong token in twenty tries disqualifies the configuration outright regardless of speed.
| config | VGPR / spills | tok/s_gen | tokens |
|---|---|---|---|
| 256 cap, G=96 (shipping) | 256 / 77 | 119–120 | 64/64 correct |
| 192 cap, G=96 | 192 / 310 | 105.6–105.8 | 64/64 correct |
| 192 cap, G=192 | 192 / 310 | 61.6–68.1 | wrong from token 14 through 25 |
The 192-cap alone cost 12% — the compiler filled the freed registers with spills, not more residency, and spills in the middle of a dot-product loop are worse than the registers they replaced. G=192 made that worse and stopped producing correct output partway through a run.
No error. No crash. Just wrong.
Nothing about that second failure looked like a bug from the outside. The process didn't crash. It didn't hang. It printed no error. It generated sixty-four tokens, same as every other run, and thirteen of them in the middle of the sequence were simply the wrong ones, with argmax reading zero — the value you get from GPU memory that was never written.
The mechanism, once you go looking for it: at G=192 there are two blocks worth of work per compute unit, but the desktop this runs on isn't a dedicated compute box — a compositor shader, an idle browser tab, whatever else has a claim on the GPU, can occupy an LDS allocation that leaves one CU able to host only one of its two assigned blocks at the moment the kernel launches. That block never starts. Its counter increment never arrives. The bounded spin on every other block eventually hits its bound anyway — because it's bounded, by design, so a genuine failure doesn't hang the GPU forever — and every block proceeds as if the barrier had been satisfied. The blocks that got real data write real values. The block that never ran writes nothing, because it never ran. The host reads the result and has no way to know one ninety-sixth of it never happened.
This is not a new lesson so much as a new and sharper instance of one this project already had a rule for: a clean, repeatable, non-crashing run is not evidence that the run measured what you think it measured. The last time that rule got written down, the failure was a request flag that silently didn't take effect. This time it's a hardware residency assumption that silently doesn't hold, on the exact class of kernel — one that assumes total occupancy for correctness, not just performance — where that assumption failing produces wrong answers instead of a slowdown.
The fix shipped today: the grid barrier now prints a fail word on every single run, pass or fail, so a zero-length silence can't hide a one-in-many failure the way it did here. tools/mega-gate.sh refuses to pass if that word ever appears. Register-cap occupancy is closed for this kernel, with the numbers to show it, not assumed away.
Where the real gains came from instead
With occupancy off the table, the rest of the day's work went into the GEMM phases and the attention kernel.
Staging activations in shared memory. Every one of the roughly 512 to 768 waves active in a GEMM phase independently re-reads the same activation row from L0/L1 cache and unpacks it from bf16 to f32 in-register. On the widest phase — the feed-forward down-projection, K=12288 — that's roughly 100 MB of repeated activation reads against 28 MB of actual weight traffic, plus the unpack cost sitting in the same instruction stream as the weight loads. Staging that row into shared memory once per block, as f32, turns every wave's read into a plain load with no unpack and no repeated global-memory traffic.
At the largest staging buffer tried — 55.6 KB of a compute unit's 64 KB shared-memory budget — this bought +6.3% (120.1 → 127.6 tok/s_gen), and also tripped the same bounded-spin grid barrier once, during testing, under real desktop LDS pressure: exactly the failure mode above, on a kernel that otherwise gated clean. A grid-barrier kernel that can be evicted from a compute unit by whatever else is running on the desktop is not a shipping configuration, full stop, regardless of its median speed. The shipped version uses a smaller, mixed f32/bf16 staging buffer — 31 KB, 1.5% slower than the aggressive version in its own stint — because the aggressive version's fail word fired once, and the mixed one's never did.
Split-K, tried and reverted. Three of the GEMM phases have a wide output dimension (N=4096 or N=12288) relative to how many waves are actually issued, which means each wave does more sequential work than the others and the rest sit idle at the barrier waiting for it. Splitting the K dimension in half so two waves share what used to be one wave's row was predicted to buy back most of that imbalance, +6%. It didn't: halving K per work-item also halves how many loads each wave has in flight before it needs the result, which cost more in lost memory-level parallelism than the imbalance it was fixing ever cost. Measured 119.0 → 108.1, a 9% loss, and reverted the same day it was tried. The shared inline dot-product body it needed along the way was kept; the split itself was not.
The attention kernel was latency-bound, not throughput-bound. Attention decode only occupies 16 of the 96 blocks in the grid — one per head — and does under a megaflop of real arithmetic per token. Its cost was never the math; it was two mostly-serial per-thread chains, one computing a score across 256 dimensions with a scalar load per element, one accumulating an output across roughly seventy sequence positions. Sharing one kernel body between the persistent kernel and the older per-launch kernel — so any fix helps both at once — widening the score loop's loads from 4-byte scalars to 32-byte vectors, and hoisting the output loop's loads eight positions at a time, took that phase from a noisy 295–605 microseconds down to a flat 92, run after run. Predicted +1.5–2.5% on the whole token; landed +4.5%, because the old phase's worst runs, not just its typical ones, are what a median absorbs, and cutting the phase's variance turned out to be worth more than cutting its mean.
The day, tallied
| stage | tok/s_gen (20-prompt median) |
|---|---|
| thirty-odd launches per token (yesterday) | 67.1 |
| one persistent launch per token | 82.0 |
| + Q4_0 weights | 115.0 |
| + LDS-staged activations | 125.4 |
| + shared, vector-widened attention kernel | 130.7 |
Five rounds, each one frozen before it ran — question, instrument, predicted range, and the condition that would falsify it, committed first. One of the five was reverted. One produced a correctness bug with no visible symptom until someone diffed the token stream against a reference, byte by byte, and is the reason every run now prints a word that says so out loud.
Full protocol files, every prediction and every falsification, in bench/mega-structural-protocol.md, bench/q4-splitk-protocol.md, and bench/attn-latency-protocol.md in amarbaro/mojo-baro.