I recently spent a whole day tracking down a numerical correctness bug in a browser-based LLM inference pipeline I built.
The root cause sat two abstraction layers below the symptom - a one-character-class mistake in an integer-bound prover inside a machine-learning compiler1 - and it only existed in the composition of four systems that each looked fine in isolation.
Here's the story of how I found the bug and how it was fixed.
Some background: running LLMs in browsers, in pieces
Modern small language models are big enough that "where" they run matters. A 360-million-parameter model in fp16 takes ~700 MB of weights, and WebGPU's buffer budget on a single browser tab has a ceiling (typically 128MB on mobile devices). One way to fit more model is to split it: compile the model into N pieces at build time, run each piece in a separate browser tab, and relay the boundary activation between tabs over WebRTC data channels.
The project I was working on (wgpu-pp) does exactly this with
SmolLM2-360M-Instruct - a small Llama-architecture model
compiled via a forked MLC-LLM pipeline-parallel pass. Two tabs each run their own "stages" (stage 0 has layers 0–15 and stage 1 contains layers 16–31), and layer actications are passed between them over a WebRTC ring.
This is my implementation of a technique is called Recurrent Pipeline Parallelism (RPP): while stage 1 is producing token A, stage 0 is already computing the boundary for token
B, so neither tab is idle.
To achieve this you need several sequences in flight at once.
If you only run one sequence at a time, stage 1 sits idle while stage 0
works and vice versa. With N sequences circulating, both shards stay
busy and throughput roughly scales. So we want numSamples > 1s.
Here's an expanded example of RPP with 4 stages. A, B, C are samples that pass through this pipeline:
Stage 0 Stage 1 Stage 2 Stage 3
┌──────┐ ┌──────┐ ┌──────┐ ┌──────┐
│S0: A │──────▶│S1: A │──────▶│S2: A │──────▶│S3: A │── argmax ──┐
│ │ │ │ │ │ │ │ │
│S0: B │ │S1: B │ │S2: B │ │S2: B │ │
│ │ │ │ │ │ │ │ │
│S0: C │ │S1: C │ │ │ │ │ │
└──┬───┘ └──────┘ └──────┘ └──────┘ │
│ ◀─────────────── token A ──────────────────────────────────┘
│
▼ next step: S0 sends B's boundary, all stages advance
Some background: the paged KV cache
An autoregressive transformer generates tokens one at a time. At each step it attends over all the keys and values it has produced so far. Recomputing those from scratch every step is wasteful, so every practical implementation caches them - the "KV cache":
step 1: attend over [K0 V0] → emit token1
step 2: attend over [K0 V0 K1 V1] → emit token2
step 3: attend over [K0 V0 K1 V1 K2 V2] → emit token3
▲
recomputed from scratch every step = wasteful
WITH a KV CACHE:
step 1: compute K0,V0 → STORE attend over cache → token1
step 2: compute K1,V1 → APPEND to cache attend over cache → token2
step 3: compute K2,V2 → APPEND to cache attend over cache → token3
(only 1 new K/V computed per step)
MLC-LLM (and vLLM, and most modern runtimes) uses a paged KV cache,
borrowed from OS virtual memory. Instead of one contiguous buffer per
sequence, the cache is a pool of fixed-size pages (here: 16 slots ×
5 key/value heads × 64 head-dim). Each sequence is described by a small
page table: an indptr (where my pages start in the global list) and
a list of page indices (which physical pages are mine). A sequence
holding 3 tokens uses part of one page; a sequence holding 50 tokens
uses 4 pages with the last one partially full.
This is great for memory efficiency (no fragmentation, easy sharing for things like prefix caching) but it means every attention kernel has to read a tiny "page table" to find its data, and every append kernel has to write to the right slot of the right page. Those tables live on the CPU and get uploaded to the GPU each step. Getting any of this wrong will cause the model to read garbage.
"The cat sat on the mat"
│
▼
┌─────────────────────┐
│ TOKENIZER │
│→ [t0 t1 t2 t3 t4 t5]│ (6 prompt tokens)
└─────────────────────┘
│
▼
┌───────────────────────────────────────────┐
│ PREFILL │
│ one parallel forward pass over ALL 6 │
│ prompt tokens at once (not one at a time)│
│ → computes K0..K5, V0..V5 per layer │
└───────────────────────────────────────────┘
│
▼
┌───────────────────────────────────────────┐
│ ALLOCATE + WRITE TO PAGED CACHE │
│ new sequence → grab a free physical page │
│ e.g. page 3 → write K0..K5,V0..V5 into │
│ slots 0-5 │
│ page_indices[seq] = [3] indptr updated │
└───────────────────────────────────────────┘
page 3: [K0 K1 K2 K3 K4 K5 __ __ __ __ __ __ __ __ __ __]
└─── prefill ───┘└─────── 10 slots free ───────┘
│
▼
┌──────────────────────────────────────────────────────────┐
│ DECODE LOOP (repeats) │
│ │
│ 1. compute K_new, V_new for the ONE new token │
│ │ │
│ ▼ │
│ 2. APPEND kernel: read page table → find next free slot │
│ → write K_new,V_new there (slot 6, 7, 8, ... of pg 3)│
│ │ │
│ ▼ │
│ 3. ATTENTION kernel: read page table → gather ALL K,V │
│ across every owned page (page 3, slots 0..k) │
│ │ │
│ ▼ │
│ 4. sample next token from the logits │
│ │ │
│ └──────────── ↺ loop: go to step 1 │
│ (until EOS or max length) │
└──────────────────────────────────────────────────────────┘
│
▼
emitted token stream
t6, t7, t8, t9, ... (output)
Some background: what a GPU kernel launch looks like
A GPU compute kernel is launched as a grid of workgroups, each
workgroup a block of threads (e.g. 256). The total number of threads is
grid × block. Each thread computes a blockId and threadId and
figures out what work to do from those:
GRID
┌───────────────────────────────────────────────┐
│ Block 0 Block 1 Block 2 ... │
│ ┌────────┐ ┌────────┐ ┌────────┐ │
│ │ 256 │ │ 256 │ │ 256 │ │
│ │threads │ │threads │ │threads │ │
│ └────────┘ └────────┘ └────────┘ │
└───────────────────────────────────────────────┘
total threads launched = grid × block
each thread computes: global_id = blockId * blockDim + threadId
Almost always, your work size isn't a clean multiple of the block size.
If you have 320 elements of work and launch 256 threads per block, you
need ceil(320 / 256) = 2 blocks = 512 threads. Threads 320–511 are
"excess" - they have nothing real to do. The standard way to handle
this is a bounds predicate: the kernel does its work inside
if (global_thread_index < work_extent) { ... }, and the excess
threads do nothing and exit.
work_extent = 320, block = 256, grid = ceil(320/256) = 2 → 512 threads launched
Block 0 (threadId 0..255) Block 1 (threadId 0..255)
global_id: 0 ─────────── 255 global_id: 256 ──────── 511
┌──────────────────────────┐ ┌──────────────────────────┐
│██████████████████████████│ │████████████░░░░░░░░░░░░░░│
└──────────────────────────┘ └──────────────────────────┘
all real work (< 320) 256..319 real 320..511 EXCESS
▲ ▲
real work threads with
(64 elems) nothing to do
█ = real work (indices 0–319) ░ = excess threads (320–511)
This is so universal that serious GPU compiler stacks (CUDA's launch
machinery, TVM's loop scheduling) insert these predicates for you. You
write a clean loop over ntoken elements; the compiler splits it into
blocks and threads and adds the guard automatically.
for each thread:
global_id = blockId * blockDim + threadId
│
▼
┌───────────────────────┐
│global_id < work_extent?
└───────────────────────┘
│yes │no
▼ ▼
┌─────────────┐ ┌─────────────┐
│ do the real │ │ do nothing │
│ work, r/w │ │ and exit │
│ memory[id] │ │ (excess) │
└─────────────┘ └─────────────┘
this "if (global_id < work_extent) { ... }" is inserted:
- by hand, or
- automatically by CUDA launch machinery / TVM loop splitting
Without this such a guard you get a specific and very nasty class of bug: excess threads read and write memory they shouldn't, in a pattern that's deterministic but looks nothing like a "normal" out-of-bounds.
normal (guarded) guard dropped (broken)
───────────────── ───────────────────────
real [0..319] → write OK real [0..319] → write OK
excess[320..511]→ skipped excess[320..511]→ writes anyway!
│
▼
reads/writes memory[320..511]
- deterministic, reproducible,
but NOT a normal OOB crash:
it silently corrupts whatever
lives just past the buffer,
every single run.
The bug, as it appeared
I wrote a baseline testing harness, which compared the outputs of an unsplit model with a split model to ensure that the unsplit models' outputs were correct. I did this as I was confident that the unsplit model's outputs would not output incorrect sequences.
Boy was I wrong.
When we turned on numSamples = 4 (four sequences in flight, all seeded
from the same beginning-of-sequence token), the output looked like this:
seq 0: [286, 311, 433, 2808, 259, 604, 1007, 1195] ← "correct"
seq 1: [ 31, 11870, 354, 12077, 198, 11870, 354, 12077]
seq 2: [ 31, 11870, 354, 12077, 198, 11870, 354, 12077]
seq 3: [ 31, 11870, 354, 12077, 198, 11870, 354, 12077]
Three things jumped out:
- Sequences 1–3 produced identical output to each other.
- Sequence 0 produced something different, which happened to match my the outputs of the unsplit, whole-model.
I hypothesizes that this was a WebGPU-async hazard: the page-table auxiliary data wasn't synced to
the GPU before the attention kernel read it. This was plausible -
WebGPU is aggressively asynchronous, and I'd been bitten by a similar
"read before the GPU finished" race before whilst working on this project. I tested this by adding an explicit
await device.sync() between the data upload and the kernel, and see if sequences 1–3 start matching sequence 0.
Nope. It changed nothing.
The first surprise: the order matters
I built a small harness that decoded one token on each of the four sequences after a fresh reset. First finding: the unsplit whole model reproduced the bug. So this wasn't a bug in my pipeline-split pass. Whatever was wrong was wrong in the TVM WebGPU runtime, in the paged KV cache, or in the kernels themselves.
Then I tried decoding the sequences in reverse order: 3, 2, 1, 0
instead of 0, 1, 2, 3.
seq 0: [9.98, 1.27, 1.71, -2.31] argmax=31
seq 1: [9.98, 1.27, 1.71, -2.31] argmax=31
seq 2: [9.98, 1.27, 1.71, -2.31] argmax=31
seq 3: [2.14, 9.34, 6.33, -3.29] argmax=286 ← now seq 3 is the odd one out
I initially thought that the decoding order did not matter - I was wrong. Reverse the order and seq 3 becomes the "correct" one. This made me realize the real invariant wasn't about sequence id at all:
The sequence whose KV lands on physical page 0 is the one that differs. Sequences on pages > 0 all agree with each other.
The paged KV cache allocates pages from a free pool, and the pool pops from the back, so the first sequence to be allocated gets page 0, the next gets page 1, and so on. The "seq 0 is correct" observation was just "whoever grabbed page 0 happens to differ" - a coincidence of allocation order, not a property of the sequence id.
This also killed my upload-sync hypothesis. The upload
machinery doesn't know or care which physical page a sequence got; if
it were a sync race, the "correct" sequence wouldn't correlate with
page assignment. And every decode in my harness already did a full
await device.sync() before reading anything back to the CPU. There
was no race to be had.
Tracing what the GPU was actually told to do
To see exactly what was happening, I monkey-patched the WebGPU API -
GPUDevice.createBuffer, GPUQueue.writeBuffer,
GPUCommandEncoder.beginComputePass,
GPUComputePass.setBindGroup/setPipeline/dispatchWorkgroups - and
recorded every buffer write (with its int32 payload) and every kernel
dispatch (with its pipeline name) during a seq-0 decode and a seq-1
decode. Then I diffed the two traces.
The diff was tiny and entirely expected:
// seq 0:
write buf#366 i32=[0] // page_indices = [0]
write buf#381 i32=[0] // append_position_map = [0]
// seq 1:
write buf#366 i32=[1] // page_indices = [1]
write buf#381 i32=[16] // append_position_map = [16]
Everything else was working exactly as expected. All the extra data uploads were there and correct. Because WebGPU’s writeBuffer
command automatically waits to finish before running any later submit() commands, we knew for sure the data would arrive before the
compute shaders tried to read it. The data upload process was definitely not the problem. This means the bug must be inside the
shader kernel code itself.
So I dumped the actual WGSL (WebGPU Shader Language) source of the
compiled kernels. The append kernel - tir_kv_cache_transpose_append - looked like this (lightly paraphrased):
// pages: (num_pages, 2, 5_heads, 16_slots, 64_dim)
let token = (block_id * 256 + thread_id) / 320; // which (token,head,dim)?
if position_map[token] != -1 { // the ONLY guard
let pos = position_map[token];
pages[pos/16, 0, head, pos%16, dim] = k_data[token, head, dim];
pages[pos/16, 1, head, pos%16, dim] = v_data[token, head, dim];
}
The only bounds protection is the -1 sentinel in position_map. If
the token index is out of range, the kernel reads position_map past
the end and either gets -1 (and skips) or gets something else (and
writes).
For ntoken = 1, the real work is 1 × 5 × 64 = 320 (5 heads, 16 slots) elements. The
kernel launches ceil(320 / 256) × 256 = 512 threads. Threads 320–511
compute token = 1 - out of range. They read position_map[1].
What's in position_map[1]? The runtime only ever writes
position_map[0] (it's a ntoken-length array, and ntoken is 1).
position_map[1] is never touched, so it sits at its zero-initialized
default. Since 0 is not -1, the guard passes, and the excess
threads write k_data[1, ...] and v_data[1, ...] - also
zero-initialized, because the kernel only wrote the real token's
data - to pages[position_map[1] / 16, ...] = pages[0, ...].
So it wrote to the same page (0), and same slot (0), on every single decode.
Reading the cache directly to confirm
I captured the live GPU buffer for layer 0's pages tensor and read it
back immediately after a seq-0 decode. All sequences share the same
input token at position 0, so every sequence's cached K/V at slot 0
should be byte-identical. Here's what page 0 actually looked like:
page 0 K[head0] = 0, 0, 0, 0, 0, 0 ← zeroed
page 0 V[head0] = 0, 0, 0, 0, 0, 0 ← zeroed
page 0 K[head3] = -0.07, -0.03, 0.40, -0.37, 0.14, 0.05 ← intact
page 1 K[head0] = 0.54, 0.32, 1.16, -1.67, 2.93, 0.57 ← intact (seq 1's data)
Heads 0, 1, 2 of page 0 / slot 0 were zero. Heads 3 and 4 were fine. And the value tensor - which has no rotary position encoding applied, so "all zeros" can't be a rotation artifact - was also zero. A bunch of zeroes written by threads that shouldn't have been writing.
Three out of five heads. Three out of five heads, because the 192
excess threads arranged themselves as exactly three groups of 64 (each
group handles one (token, head) pair): out-of-bounds token 1, for
heads 0, 1, 2. Head 3 and 4 weren't touched because the next three
groups (heads 3, 4, and a fifth out-of-bounds one) compute token = 1
too, but their position_map reads also come back zero and they write
to the same locations - harmless overwrites of the same zeros.
Here's a diagram of what this looks like:
Why this bug was difficult to spot
Here's the part that made this bug difficult to spot.
My baseline testing harness compared the 2-stage split pipeline against an
unsplit whole-model run (batch_decode, all 32 layers in one
instance). The split and the unsplit matched each others' outputs - so I trusted the
unsplit output as ground truth.
But the unsplit model is also single-sequence. It only ever uses page 0. And page 0 was being corrupted on every decode. So the outputs of the unsplit model (the "oracle") was corrupted exactly the same way as the thing it was checking. They matched because both sides were broken identically.
What I believed: What was actually true:
seq 0 = CORRECT seq 0 = CORRUPTED (heads 0–2 zeroed)
seqs 1–3 = WRONG seqs 1–3 = CORRECT (pages 1–3 intact)
oracle = GROUND TRUTH oracle = CORRUPTED (single-sequence, page 0)
The tokens I'd flagged as "wrong" - [31, 11870, 354, 12077, ...] -
were the real greedy continuation. The tokens I'd flagged as "correct" -
[286, 311, 433, 2808, ...] - were the output of a model with two
attention heads silently zeroed on every layer. The bug was making the
truth look like the defect.
I confirmed this after the fix: a real prompt fed token-by-token - "The capital of France is" - produces " Paris.\n" matching HuggingFace fp32 (32-bit floating-point weights) before diverging into coherent text, (this was expected because the model was 4-bit quantized - uses less memeory but its output is generally less accurate than the original model). Same result on seq 0 and seq 1. The new "true" sequences were what the "wrong" sequences had been saying all along.
Why the guard was missing in the first place
Okay so why was there no guard/check in place for this?
Buckle up because this part involves quite a bit of math.
Personally, this was toughest part of my investigation, as it required a lot of
thought and checking by hand. Fortunately, it does end up paying off in the end :).
The append kernel is written as a clean three-deep loop over
ntoken × num_kv_heads × head_dim. It gets scheduled into a GPU kernel
by TVM's dlight gpu.Fallback rule, which does:
schn = sch.fuse(spatial_loops) # → one loop of extent ntoken*320
bx, tx = sch.split(schn, [None, 256]) # → ceil(n*320/256) blocks of 256
sch.bind(bx, "blockIdx.x")
sch.bind(tx, "threadIdx.x")
split is supposed to add a T.where(bx*256+tx < ntoken*320) predicate
when the split is imperfect. It skips the predicate if it can already
prove it's true: if (!CanProve(predicate, kSymbolicBound)) add_predicate(...). For a fused dynamic extent like ntoken*320, it
"proved" the predicate was always true.
Well, turns out this proof was wrong.
The proof goes like this. We want to show bx*256 + tx < n*320 given
bx < ceil(n*320/256) and tx < 256. Rearranging, that's equivalent to
n*320 - bx*256 - tx ≥ 1. The minimum of bx*256+tx over the valid
range is (ceil(n*320/256) - 1) * 256 + 255 = ceil(n*320/256)*256 - 1.
So the minimum of the difference is:
min = n*320 - (ceil(n*320/256)*256 - 1)
= n*320 - ((n*320 + 255) // 256) * 256 + 1
= (n*320 + 255) % 256 - 254
The analyzer asks const_int_bound for the range of
(n*320 + 255) % 256. There's a "tighter bound via modular set"
optimization that uses what it knows about the dividend's residue
class. For n*320 + 255, the modular-set analysis says "this is
congruent to 255 mod 320" (coeff = 320, base = 255). The bound code
then needs to find the smallest non-negative value congruent to 255
mod gcd(320, 256) = 64. The correct residue is 255 % 64 = 63, and
the possible values mod 256 are {63, 127, 191, 255}, so the correct
bound is [63, 255].
But the code did this:
int64_t base_mod = mod_a->base % modulus; // 255 % 256 = 255 ← WRONG
if (base_mod < 0) base_mod += modulus;
int64_t tight_max = modulus - gcd_coeff_mod + base_mod; // 256 - 64 + 255 = 447
if (tight_max >= modulus) tight_max -= modulus; // 447 - 256 = 191
return MakeBound(base_mod, tight_max); // [255, 191] ← min > max!
It normalized the base by % modulus (256) instead of
% gcd(coeff, modulus) (64). That gives base_mod = 255, then it
tries to paper over the resulting out-of-range tight_max with the
if (tight_max >= modulus) tight_max -= modulus line - which produces
[255, 191], a bound where min is greater than max.
Subtract 254 and you get [1, -63], a bound where min (1) is still ≥
the threshold I needed (≥ 1). So CanProve returns true. The
predicate is treated as redundant. split drops it. The kernel ships
unguarded, and 192 excess threads go on a writing spree on every
single decode.
I confirmed this with a fifteen-line test:
# sch.fuse(3 dynamic loops) → sch.split([None, 256])
# BEFORE the fix:
for bx, tx in grid(ceil(n*320/256), 256):
with sblock: ... # no T.where - guard lost
# AFTER the fix:
for bx, tx in grid(ceil(n*320/256), 256):
with sblock:
T.where(bx*256 + tx < n*320) # guard restored
...
Splitting a plain n-extent loop always had the predicate; the bug
only triggered after fuse produced a dynamic extent, which
is why it escaped notice - most test cases don't fuse before splitting a
symbolic-bound loop.
This affected every Fallback-scheduled kernel with a fused dynamic extent, not just the append kernel. The append kernel was the one that wrote to persistent state.
The fix
The fix involved modifiying the truncated-Mod and FloorMod ops in
tvm-src/src/arith/const_int_bound.cc. The change reduces the base by
the GCD instead of the modulus:
// BEFORE - wrong: base normalized to [0, modulus)
int64_t base_mod = mod_a->base % modulus;
if (base_mod < 0) base_mod += modulus;
int64_t tight_max = modulus - gcd_coeff_mod + base_mod;
if (tight_max >= modulus) tight_max -= modulus; // papering over the wrong base
return MakeBound(base_mod, tight_max); // → [255, 191]
// AFTER - correct: base normalized to [0, gcd)
int64_t base_mod = mod_a->base % gcd_coeff_mod;
if (base_mod < 0) base_mod += gcd_coeff_mod;
int64_t tight_max = modulus - gcd_coeff_mod + base_mod;
return MakeBound(base_mod, tight_max); // → [63, 255]
The truncated-Mod also needs to handle a negative dividend's
mirrored residue (truncated mod can go negative; floor mod always lands
in [0, modulus) so it's simpler). I then rebuilt the libtvm_compiler,
and recompiled both model wasms (split, and unsplit model).
The shipped tir_kv_cache_transpose_append_kernel WGSL now contains
the guard:
if (((block_id * 256) + thread_id) < (podArgs.ntoken * 320)) {
pages[...] = k_data[...];
}
Which is the guard we wanted.
Verification
After the fix:
- All 4 sequences produce identical output, in any decode order, on both the split and the unsplit model.
- The shared output is the former "wrong" one - the old "correct" tokens were the corrupted party, exactly as the page-0 corruption predicted.
- A real prompt ("The capital of France is") fed token-by-token produces " Paris.\n", identically on seq 0 and seq 1.
The lessons
Two things made this bug hard.
The corruption was deterministic and uniform across the "wrong" sequences. That's the exact shape you'd expect from a data/indexing bug, which I thought was the case here. It wasn't. The uniformity came from a totally different mechanism (excess threads all reading the same zero sentinel and writing to the same page 0), but the shape of the symptom was misleading. "Deterministic and uniform" is not actually strong evidence for "simple indexing bug"; it's just evidence that the bug doesn't depend on timing.
The reference was corrupted. I'd built a testing harness
specifically to catch this class of regression, and I assumed all was well for some time.
But the outputs of the unsplit model (oracle) in the harness was itself single-sequence - the exact
configuration where the corruption is always present - so it was
confirming the bug instead of catching it. The only way out was to
stop trusting the oracle and measure the hardware state directly
(read the pages buffer back from the GPU and look at it).
And the root cause lived two abstraction layers below where the
symptom appeared: a model-output bug caused by a kernel-write bug
caused by a missing guard caused by a wrong proof caused by a
modular-arithmetic bound that returned min > max. Each layer looked
fine in isolation - the kernel was a reasonable kernel, the scheduler
was a reasonable scheduler, the prover was a reasonable prover, the
bound code was a reasonable bound code with a subtle off-by-GCD error.
The bug only existed in the composition.
This is the kind of bug that's worth writing up not just because the fix was hard (a one line change deep down in the software stack) but because the finding was hard, and because the tools that found it - direct GPU buffer readback, monkey-patched API tracing, a willingness to distrust the oracle - are generally useful.
The next time a correctness bug looks "deterministic and uniform," and the next time a trusted reference turns out to share the bug it's supposed to catch, I'll know to go straight to the hardware.
1You can view PRs for my fixes here: https://github.com/apache/tvm/pull/19978