Phi-3-mini (3.8B) running in the browser on 10 hand-written WGSL kernels — measured +16% on total wall-clock throughput and +31% on decode against WebLLM on an Apple M2 Max, with no TVM, no compiler, and no WASM runtime.
| WebLLM (TVM) | Zero-TVM (this repo) | |
|---|---|---|
| Total throughput, prefill + decode (M2 Max, same weights, same session) | 60.0 tok/s | 69.55 tok/s (+16.0%, 2026-07-30 — see BENCH.md) |
| Decode only, same pair | 63.2 tok/s (self-reported) | 83.1 tok/s (+31.4%) |
| Time to first token, ~35-token prompt | ~150 ms (implied) | 291 ms — WebLLM wins this one; see below |
| Unique WGSL kernels | 85 (autotuned) | 10 roles / 55 kernels (37 .wgsl + 18 generated int4 variants) |
| Total WGSL lines | 12,962 (generated) | ~2,150 hand-written + a 280-line readable generator |
| Dispatches per decode step | 342 | 228 (f16 KV) / 260 (int8 KV) |
| Runtime | TVM → WASM scheduler | Plain TypeScript, none |
| Tokenizer | bundled from WebLLM | BPE from scratch (tokenizer.ts) |
| JS bundle (chat page, excl. weights) | 5.9 MB / 2.1 MB gz | 157 kB / 33 kB gz |
Same model, same quantized weights. WebLLM / MLC-LLM — the standard way to run a browser LLM — ships an Apache-TVM pipeline that emits 85 autotuned WGSL kernels driven from a WASM scheduler. This repo replaces that whole stack with 10 kernel roles (55 WGSL kernels — 37 hand-written files plus 18 int4-matmul variants emitted by a small readable generator, counting subgroup/tiled/int8 variants) and ~2,000 lines of TypeScript (engine + tokenizer + weight loader) — and, as of the 2026-07-30 head-to-head, runs faster than it on the one machine measured. The whole forward pass — 32 transformer layers, paged KV cache, int4-dequant matmul, RoPE, fused FFN, RMSNorm, paged attention, argmax sampling — is readable end-to-end in a single sitting. That is the point.
Two numbers, always both. Total is wall-clock throughput with prefill included — both engines doing identical work, the conservative figure. Decode excludes prefill and is the kernel-level figure. Quoting one without the other is cherry-picking, so this repo quotes both everywhere. (Until 2026-07-30 it quoted one blended number labelled "decode"; the bench harness had also stopped making our half pay prefill after cross-turn prefix reuse shipped, which invalidated two published Qwen comparisons. The defect, the fix, and the withdrawn pairs are written up at the top of BENCH.md.)
Numbers above are measured from the source and build output in this repo; throughput varies by GPU, shows live in the chat UI, and the full head-to-head (methodology + raw runs) is in BENCH.md. Zero-TVM issues fewer dispatches than TVM because it fuses operations TVM's default pipeline doesn't:
qkv_fused.wgsl— Q/K/V projection + RoPE + paged-KV append, one dispatch per layer (was 3 in earlier revisions and in TVM's emission).attention.wgsl— paged attention combined with the page-table read.fused_ffn.wgsl— gate + up + SiLU in one pass.add_norm.wgsl— residual add + RMSNorm in one pass.
Counting note: "10 kernel roles" counts the default f16-KV decode path. By file
basename there are 11 prefixes — the eleventh, kv_quantize_int8.wgsl, only
runs on the opt-in ?kv8=1 int8-KV path.
Every FLOP the model executes is in a file you can open. Every GPU buffer has a human label. Every dispatch is annotated in src/zero-tvm/engine-core.ts — the single decode engine that both the chat page and the validation page drive.
Hand-written GPU kernels usually lose significantly to an autotuning compiler. The claim this repo is designed to test is: for a decoder-only LLM of this shape, most of the compiler's complexity budget isn't buying much. The expensive parts are matmul, attention, and int4 dequant. Everything else is plumbing. ~10 kernels of plumbing, instead of 85.
Measured result on an Apple M2 Max, Chrome 150, identical Phi-3-mini-q4f16_1 weights, same-session head-to-head (2026-07-30): Zero-TVM is 16.0% faster on total wall-clock throughput and 31.4% faster on decode than the autotuned compiler. Absolute tok/s drifts a lot with machine state — both engines moved together by ~10 tok/s between 2026-07-29 and 2026-07-30 with no code change on this path — so only same-session pairs mean anything; the current medians are in the table above and in BENCH.md, synced from bench/results.json. Full methodology and A/B results in BENCH.md, including three tile-variant experiments, a prompt-lookup speculative-decoding experiment, and an FFN-prologue-fusion experiment that were falsified by measurement rather than shipped. (An earlier measurement on an M2 Pro with an older, buggier engine put this repo 22% behind WebLLM — that history is preserved, not retconned, in BENCH.md.)
The advantage grows with how recent the architecture is — the most interesting thing in the corrected round, and monotonic on both metrics:
| Model | arch. released | Δ total | Δ decode |
|---|---|---|---|
| Phi-3-mini | 2024 | +16.0% | +31.4% |
| Qwen3-4B | 2025 | +31.7% | +58.0% |
| Qwen3.5-4B | 2026 | +100.5% | +113.6% |
The reading that fits is that compiler stacks have had less time to tune newer architectures, so there is more room for a hand-written kernel set to take. That is an observation across three points on one machine on one day, not a proven law — no mechanism was isolated and nothing controls for how differently each model stresses the two engines. Treat it as a hypothesis worth testing on a fourth model.
Where we lose: time to first token on short prompts. WebLLM's prefill runs 251–271 tok/s on Phi-3/Qwen3, implying a TTFT around 150 ms on Phi-3 where we measure 291 ms — and 453 ms on Qwen3-4B. We win sustained decode decisively and lose the first-token sprint on short inputs. It is specifically a short-prompt weakness: chunked prefill measures 202 tok/s on an 816-token prompt, and cross-turn prefix reuse removes prefill entirely on follow-up turns. It is the top open item on BENCH.md's levers list.
The scope of these numbers matters: one machine, one browser, short-context bench, one same-session pair per model. But within that scope the question the repo was built to test now has a sharper answer — the ten hand-written kernels aren't merely "close enough" to the 85 autotuned ones; on this hardware they win on throughput. The repo also makes the stack auditable: if you want to instrument a layer, add a new fusion, test a different attention pattern, or teach someone how browser LLM inference works at the metal, there is no compiler in the way — just WGSL and a few hundred lines of TypeScript orchestrating it.
The closest reference point is Karpathy's llm.c (hand-written CUDA/C GPT-2). This is that thesis — you don't need the giant framework — ported to browser / WebGPU / int4 / paged KV / modern arch, for a model people actually use.
Requirements: A recent Chrome or Edge with WebGPU enabled and the shader-f16 feature available. Tested on macOS (M2 Pro / M2 Max, Chrome 120–150). Other platforms should work but are untested.
npm install
npm run devThen open http://localhost:5173/zero-tvm.html. On first load the browser downloads the Phi-3-mini-4k-instruct Q4 weights from HuggingFace (~1.8 GB) and caches them in OPFS (Origin Private File System). Subsequent loads are instant.
To build a deployable bundle:
npm run build # → dist/The build produces a multi-page Vite output: index.html (landing page — project overview, shader catalog, compare table), zero-tvm.html (chat demo), compiler-chat.html, demo.html (dispatch visualization), validate.html (multi-prompt smoke test), webllm-bench.html (head-to-head harness), architecture.html, docs.html, dump.html (TVM shader capture), shaders.html (captured-shader browser).
zero-tvm.html accepts a handful of query flags for A/B-ing shader variants without rebuilding. Defaults are tuned for Apple GPUs; flags let you isolate a path:
?sg=0— disable all subgroup shaders (argmax / attention / QKV matmul)?sgqkv=0/?sgattn=0/?sgargmax=0— disable one at a time?qkvtile=1/?qkvtile2=1— opt into tiled QKV variants (both regressed on M2 Pro — kept for portability testing)?ffnsg=1— opt into the tiled-subgroup fused FFN?kv8=1— opt into the int8 KV cache path (kv_quantize_int8+attention_int8)?vec4=0/?vec4qkv=0— opt OUT of the vec4 weight + activation loads in the int4 matmuls /qkv_fused. Default ON since 2026-07-25: measured +7.1% together on M2 Max (see BENCH.md)?vec4h=0— opt OUT of just the K%512_vec4hhalf-unroll matmul siblings (the K%1024 full-vec4 instances stay on). Default ON since 2026-07-29: +5.7% on Qwen3, +2.0% on Qwen3.5 (they're what give d=2560 / ffn=9728 wide loads)?fuseqk=0— opt OUT of the fused qk_norm+RoPE+KV-append kernel on the Qwen3 unfused path (restores the 3-dispatch reference chain, 8 → 10 dispatches/layer). Default ON since 2026-07-29: +2.3% alone on M2 Max?splitk=N— split-K flash-decode attention, N ∈ 2..16 partitions per head + combine pass; f16 KV only. Default N=8 since the long-context A/B (+3.1% at 128-token, +4.0% at 1024-token decode on M2 Max);?splitk=0to disable?fuseprologue=1— fold the FFN-entry add_norm into the FFN kernel's prologue,add3_normat the layer tail. Measured −13.7% on M2 Max — a documented negative result, kept for A/B on other GPUs
Open DevTools and window.specSim(160, 3, 3) runs the CPU-side prompt-lookup speculative-decoding acceptance simulator over three prompt types — see src/zero-tvm/spec-sim.ts.
Phi-3-mini-4k-instruct (3.8B, q4f16_1) — the default. Everything above — the headline numbers, the kernel counts, the fusion story — is about Phi-3. No URL flag needed; all existing URLs keep their exact behavior.
Qwen3-4B (q4f16_1) — ?model=qwen3 on zero-tvm.html and
validate.html. A v1 port that exercises the spec-parameterized engine on
a genuinely different architecture:
- GQA 32/8 — 32 query heads over 8 KV heads (Phi-3 is 32/32 MHA), with
qDim = 4096 ≠ d = 2560 - QK-norm — per-head RMSNorm on Q and K between the projection and RoPE
- byte-level BPE tokenizer (Qwen2-style
tokenizer.json; Phi-3 is SPM) - tied lm_head — logits reuse the quantized embedding matrix
- ChatML template, run in the non-thinking form (
<think>rendering is not built)
Honest performance framing: QK-norm is incompatible with the fused
QKV+RoPE+KV-append kernel, so the QKV matmul stays a separate dispatch — but
since the 2026-07-29 tuning round everything after it is fused: the
qk_norm_rope_append kernel folds the per-head norm, RoPE, and the paged KV
write into one pass (8 dispatches/layer vs the Phi-3 chat path's 7;
?fuseqk=0 restores the 10-dispatch reference chain), and the _vec4h
K%512 matmul variants extend the wide loads to d=2560 and ffn=9728
(?vec4h=0 opts out; K=4096 instances were already on full vec4). No
int8-KV. Measured same-session pair on an Apple M2 Max (2026-07-30,
Chrome 150, identical local weight bytes, corrected protocol): Zero-TVM
59.85 tok/s total (TTFT 453 ms, decode 75.49) vs WebLLM 0.2.84's
prebuilt Qwen3-4B at 45.46 tok/s total (self-reported decode 47.77) —
+31.7% total, +58.0% decode. On this model WebLLM has the better first
token: its 263–271 tok/s prefill implies ~150 ms against our 453 ms.
Read gaps cautiously: one pair, one machine. Two earlier published pairs are
gone — the 2026-07-29 "75.7 vs 43.8, +73%" pair is withdrawn (the bench
harness had stopped making our half pay prefill; see BENCH.md's corrected
protocol), and the 2026-07-28 v1 pair (25.4 vs 14.2) was a degraded session
that did not reproduce. The Zero-TVM-vs-Zero-TVM per-item deltas from the
tuning round are unaffected and stand: fused-qk +2.3%, vec4h +5.7%, combined
+5.8% over the same-day flags-off half. Protocol and caveats in
BENCH.md; machine-readable pair in
bench/results/qwen3-4b.json.
Weights: node scripts/download-weights.mjs --model qwen3 primes the local
dev mirror (~2.3 GB); without it the page streams from HuggingFace.
Qwen3.5-4B (q4f16_1) — ?model=qwen35 on zero-tvm.html and
validate.html. The third model, and the first hybrid: 24
gated-DeltaNet (linear-attention) layers interleaved with 8 gated
full-attention layers, running on the same hand-written kernels. To our
knowledge this is the first hand-written-kernel int4 implementation of a
gated-DeltaNet hybrid running in a browser. What it adds over Qwen3-4B:
- hybrid layer stack — attention on every 4th layer only; the other 24 layers run gated DeltaNet: a recurrent delta-rule state update (16 k-heads, 32 v-heads, head dims 128, short conv K=4), no KV cache on those layers
- GQA 16/4 at head_dim 256 on the attention layers, with partial RoPE (only 64 of 256 dims rotate) and a sigmoid output gate per head
- 248,320-token vocab with renumbered special tokens (the repo's shipped
mlc-chat-config.jsonstill lists the stale Qwen3 stop ids — we resolve stops fromtokenizer.jsoninstead), tied lm_head
Measured same-session pair on an Apple M2 Max (2026-07-30, Chrome 150, identical local weight bytes, corrected protocol): Zero-TVM 65.28 tok/s total (TTFT 171 ms, decode 73.30) vs WebLLM 0.2.84's prebuilt Qwen3.5-4B at 32.56 tok/s total (self-reported decode 34.32) — +100.5% total, +113.6% decode. This is the one model where first-token latency is roughly a wash rather than a loss (our 171 ms against an implied ~0.2 s from WebLLM's 175–177 tok/s prefill).
Prefill is no longer token-by-token: the 2026-07-29 prefill round
runs prompts in chunks of ≤64 (batched projections, one gdn_recur
dispatch per layer per chunk, causal batched attention) — 67.9 → 202
prefill tok/s on an 816-token prompt — and every model now does
cross-turn prefix reuse (turn 3 of a ~950-token conversation:
first-token latency 14.3 s → 0.19 s on Qwen3.5; reused-prefix logits
verified bit-identical to a fresh prefill). The 2026-07-29 Qwen3 tuning
round's _vec4h matmuls also engage on the hybrid's K=2560 projections
(+2.0% in a Zero-TVM-vs-Zero-TVM A/B). Superseded history: the "65.7 vs 34.0,
+93%" cross-check published on 2026-07-29 is withdrawn (bench-harness
defect — our half had stopped paying prefill; see BENCH.md), and the earlier
53.07/32.36 (+64%) and 47.99/31.99 (+50%) pairs were fair but are older.
Same honest framing as always: one machine, one pair, and the GDN decode
kernels are still scalar (non-subgroup) — the decode number remains a floor,
not a tuned result. Protocol and caveats in BENCH.md;
machine-readable pair in
bench/results/qwen35-4b.json.
Weights: node scripts/download-weights.mjs --model qwen35 primes the local
dev mirror (~2.6 GB); without it the page streams from HuggingFace.
The directory layout is the narrative arc of the project. Each page is a milestone.
index.html (landing page — essay, shader catalog, compare table)
compiler-chat.html → src/compiler/chat-v2.ts (1) WebLLM reference: captures
dispatches, our shaders replay
279 of 342 of them
zero-tvm.html → src/zero-tvm/chat.ts (2) The result: all dispatches
replaced, WebLLM never touched
validate.html → src/zero-tvm/validate.ts Multi-prompt smoke test driving
src/zero-tvm/engine-core.ts
webllm-bench.html → src/webllm-bench/main.ts (3) Honesty check: WebLLM driven
against the same local weights
for a fair head-to-head
demo.html → src/demo.ts Dispatch timeline visualization
dump.html → src/dump-tvm.ts Captures all 85 TVM-emitted WGSL
shaders.html → src/dump-shaders.ts Browses the captured shaders
test-shaders.html → src/compiler/test-harness.ts Per-shader correctness vs TVM
test-chain.html → src/compiler/test-chain.ts
src/
zero-tvm/ THE RESULT
engine-core.ts ~1,000 lines — THE decode engine: buildDecodeEngine,
allocKVPages/allocKVPagesInt8, the 32-layer decode
loop. No DOM. Parameterized by mode: unfused
reference path (validate, 9 dispatches/layer) or
fused QKV+RoPE+KV-append (chat, 7/layer; 8 with int8
KV), plus two generate styles — blocking
generate/forwardLogits and the pipelined readback
ring (generatePipelined). Driven by BOTH chat.ts
and validate.ts.
chat.ts ~500 lines — thin chat page: DOM state, boot wiring
(via loading-ui's bootEngine), streaming render.
variants.ts ~160 lines — URL-flag A/B harness (?sg/?matmul=/
?kv8=1 …) and variant→pipeline resolution.
markdown.ts ~170 lines — minimal streaming Markdown renderer.
bench-console.ts ~160 lines — window.bench / benchBatched / specSim
devtools harnesses for the chat page.
spec-sim.ts 120 lines — CPU-side prompt-lookup speculative-decoding
acceptance simulator. Used to falsify a speed-up
experiment before building shaders.
tokenizer.ts ~280 lines — BPE tokenizer from scratch
weight-loader.ts ~300 lines — direct HuggingFace Phi-3-MLC fetch,
OPFS cache, layer-ordered streaming
validate.ts ~320 lines — multi-prompt forward-pass smoke test
loading-ui.ts ~280 lines — shared progress-bar UI + bootEngine
flow used by both chat and validate
webllm-bench/
main.ts Head-to-head harness: WebLLM v0.2.80 wired against
/local-weights/* so the comparison runs on identical
bits. See BENCH.md.
compiler/ THE SHADERS
compiler.ts ~280 lines — pipeline creation, weight buffer
allocation. Not an optimizing compiler — the name is
historical.
shader-prelude.ts ~70 lines — PHI3 model constants (single source of
truth) rendered as a WGSL `const` prelude that is
injected into every shader at module creation, so no
model-shape literal lives inside the WGSL itself.
shaders/ 18 hand-written WGSL files (~2,150 lines) + one
generator for the 9-variant int4_matmul family:
add_norm.wgsl Residual add + RMSNorm fused
embedding.wgsl
rms_norm.wgsl
rope.wgsl (legacy, subsumed by qkv_fused)
kv_append.wgsl (legacy, subsumed by qkv_fused)
kv_quantize_int8.wgsl int8-KV opt-in path
qkv_fused.wgsl Q/K/V proj + RoPE + paged-KV append, 1 dispatch/layer
qkv_fused_sg.wgsl subgroup-reduce variant (default on Apple)
qkv_fused_scratch.wgsl int8-KV-compatible variant (writes full V to scratch)
qkv_fused_tiled_sg.wgsl experimental tile variant (regressed — kept for A/B)
qkv_fused_tiled2sg.wgsl experimental 2-subgroup tile variant (regressed)
attention.wgsl Paged attention (vLLM-style page table)
attention_sg.wgsl subgroup-reduce variant (default on Apple)
attention_int8.wgsl int8-KV opt-in path
fused_ffn.wgsl Gate + up + SiLU, fused
fused_ffn_tiled_sg.wgsl tile + subgroup variant
int4_matmul.gen.ts generator for the int4 matmul family — the 9
former files differed only on {f16|f32 out} ×
{1|4|8 rows/WG} × {tree|subgroup reduce} × {M=1|4};
entry names are unchanged (int4_matmul, _sg,
_tiled, _tiled8, _f32, _f32_sg, _f32_tiled,
_f32_tiled8, _batched_m4). Every variant is
still plain readable WGSL — dump them all with:
node -e "import('./src/compiler/shaders/int4_matmul.gen.ts').then(m => console.log(m.debugDumpAll()))"
argmax.wgsl
argmax_sg.wgsl subgroup variant
tvm-shaders/ THE EVIDENCE — all 85 TVM-emitted WGSL kernels,
captured from a running WebLLM session by
src/dump-tvm.ts. Keep this next to compiler/shaders/
and the replacement is auditable.
RESEARCH.md is the writeup of how the shader capture worked and what reading TVM's output revealed about its kernel set. BENCH.md records the measured numbers, the head-to-head methodology, and the experiments that were falsified rather than shipped. RESEARCH_STANDARDS.md is the 15-principle engineering discipline this repo shares with its sibling WebGPU/WGSL research projects (webgpu-q quantum chemistry, webgpu-dna radiobiology, neuropulse LLM visualization) — single source of truth, falsifiable JSON artifacts, honest negatives, no fudge factors, shader byte-hashing, multi-level correctness.
Three layers, intentionally separate:
-
Per-shader correctness vs TVM —
test-shaders.html(src/compiler/test-harness.ts). Loads WebLLM, intercepts the WGSL device to capture every TVM dispatch from a real decode step, then runs each of our shaders against the matching TVM dispatch's input buffers and compares the f16/f32 output buffers element-wise (cmpF16/cmpF32). Each shader is reported withmaxDiff,avgDiff, and exact-match percentage. Catches kernel-level bugs but only exercises one prompt. -
Live forward pass on diverse prompts —
validate.html(src/zero-tvm/validate.ts). Drivesengine-core.ts(the reference path) against a battery of prompts (factual recall, arithmetic, code, instruction following, open-ended) throughforwardLogitsand reports for each: the top-10 next-token candidates with probabilities, the entropy of the predictive distribution, and a short greedy continuation. A reader can scroll through the page and verify the model behaves like Phi-3 on inputs that were never in the per-shader test set. -
Automated kernel correctness (headless, CI-ready) —
npm run test:kernels(tests/kernels/). Runs the real WGSL kernels fromsrc/compiler/shaders/against independent CPU references — 8 of the 10 roles, exact or within f16 tolerance. Uses the Dawn-native WebGPU binding on Mesa lavapipe, so it needs no GPU and runs in CI; on a machine with a GPU it uses the real adapter. This is the automated net the per-shader browser harness (layer 1) never had.
zero-tvm.html and validate.html share the same decode engine
(engine-core.ts); the chat page runs it in the fused / variant-selected /
pipelined configuration while validate runs the unfused scalar reference with
blocking per-token readback, so a validate pass exercises the same kernels,
buffers, and bind-group layout the chat page ships.
Measured on Apple M2 Max, Chrome 150, Phi-3-mini-4k-instruct q4f16_1,
128-token target × 5 runs, median, same-session pair (2026-07-30; latest run
recorded in bench/results.json). Every run pays a full prefill on both
sides:
| total tok/s (median) | |
|---|---|
| Zero-TVM (this repo, f16 KV, default shaders) | 69.55 |
| WebLLM v0.2.80 (MLC-LLM, same weights, same session) | 60.0 |
Both metrics for that pair: total 69.55 vs 59.95 (+16.0%), decode-only
83.10 vs 63.23 (+31.4%), our TTFT 291 ms against WebLLM's implied ~150 ms.
The marker-wrapped cells are the total wall-clock medians, auto-synced from
bench/results.json; the unambiguous per-model record is in
bench/results/.
Reproduce on any WebGPU GPU with npm run bench — it drives both engines on
identical weights and refreshes every marker-wrapped number (see below). It
also runs headless on a cloud GPU (Colab notebook + Docker image in
bench/) for machines without a local one.
bench/results.json is the single source of truth. npm run bench writes it
and rewrites every marker-wrapped number (<!--bench:zt/webllm/gap-->) in
BENCH.md, README.md, index.html, and docs.html, plus the // bench:zt constant
in src/webllm-bench/main.ts — npm run bench:sync -- --write re-propagates
it without a GPU (see bench/README.md). Merging to main then auto-deploys
zerotvm.com (Cloudflare Pages) and mirrors the built site + Space card to the
Hugging Face Space via
.github/workflows/deploy-space.yml (needs an HF_TOKEN repo secret; the job
skips with a notice if it's absent).
That's +16.0% on total and +31.4% on decode against the autotuned compiler on an identical workload — one machine, one model, one browser, one pair. An earlier M2 Pro measurement with an older engine read 22% behind; both hardware and code changed since, so the same-session WebLLM number is the only valid comparator (details in BENCH.md). See BENCH.md for the full protocol, the raw numbers, the 2026-07-30 harness defect and the two pairs it invalidated, and the optimization experiments that were measured and dropped:
- Three QKV tiling strategies (1152 WGs × 32 threads, 2304 × 32, 2304 × 64) — all regressed vs the 4608-WG subgroup baseline. Apple GPUs want high WG occupancy on this kernel; tiling reduces it.
- Prompt-lookup speculative decoding — CPU-simulated over three prompt types
(prose, code, summary). Acceptance rate <8% at N=3, K=3; below the 67% threshold
the
(1+αK)/((K+1)/2)throughput formula needs to break even. Falsified before any shader was written (src/zero-tvm/spec-sim.ts). - FFN prologue fusion (
?fuseprologue=1) — folding addNorm1 into the FFN kernel to remove 32 dispatch bubbles per token measured −13.7% on M2 Max: the redundant per-workgroup RMSNorm recompute costs more than the dispatch bubbles save. Kept behind its flag for A/B on other GPUs.
These are the caveats that survive the code as-shipped. Several earlier ones — silent context overflow, per-token uniform buffer leaks, double queue.submit(), redundant first-token readback — were turned into code fixes rather than documentation. The remaining items are either inherent to the approach or deliberate scoping decisions.
- Phi-3-mini-4k-instruct Q4 only. The model shape lives in one place — the
PHI3object insrc/compiler/shader-prelude.ts(D=3072,HEADS=32,HEAD_DIM=96,LAYERS=32,FFN=8192,VOCAB=32064,PAGE_SIZE=16,MAX_PAGES=257) — which is rendered as a WGSLconstprelude and injected into every shader, so the address arithmetic uses named constants (D,HEAD_DIM,KV_PAGE_STRIDE, …) instead of magic literals. Porting to Mistral or Llama is still not just a config edit: the kernels bake in structural assumptions (MHA with no GQA,HEAD_DIM = 32×3thread partitioning in attention,Ddivisible by 256, RoPE pair layout), plus the tokenizer/template/loader items below. - GPU memory footprint ≈ 3.6 GB. Phi-3-mini Q4 weights are ~1.8 GB, and the paged KV cache is
32 layers × 257 pages × 196,608 B/page ≈ 1.6 GB. On an M2 Pro with 16 GB unified memory this is invisible; on a 4 GB integrated GPU it will OOM during KV allocation before the first token. If you want to trade context length for memory, lowerMAX_PAGESinsrc/compiler/compiler.ts— 128 pages = 2048-token context, ~0.8 GB KV, which fits almost anywhere. The optional?kv8=1int8 KV path roughly halves the KV footprint. - Requires the
shader-f16WebGPU feature. Matmuls run in f16 (enable f16at the top of every int4 matmul variant). The LM head uses an f32 output buffer (theint4_matmul_f32generator variant) because the sampling pipeline needs f32 logits — TVM'sNT_matmul14_cast2does the same cast. Chrome/Edge with WebGPU andshader-f16is required; Safari's WebGPU does not yet exposeshader-f16. - BPE tokenizer is a hand-rolled reimplementation, not
tokenizers.js.src/zero-tvm/tokenizer.tsis ~280 lines: vocab lookup, merge table, metaspace prefixing, byte fallback, SentencePiece hex-byte decode. It does not implement HuggingFace's full pre-tokenization regex pipeline or Unicode NFKC normalization. For normal English chat it matches the reference tokenizer; for emoji, unusual Unicode, or some punctuation patterns it may diverge. If correctness matters for your input, run the prompt through@huggingface/tokenizersand compare. - Phi-3 chat template is baked in.
buildChatPromptintokenizer.tsemits<|system|>...<|end|>\n<|user|>...<|end|>\n<|assistant|>\n. Stop tokens are the Phi-3 set{2, 32000, 32007}. Port to another model → edit both. - Weight loader expects MLC's Q4f16_1 layout.
mlc-ai/Phi-3-mini-4k-instruct-q4f16_1-MLC, including MLC's renamed parameter scheme (transformer.h.N.mixer.*, notmodel.layers.N.*). If MLC re-quantizes or re-names, the loader needs a patch.
- Greedy decoding only. Sampling is a single
argmax.wgsldispatch. No temperature, top-k, top-p, repetition penalty. A CPU-side sampler over the f32 logit buffer would be ~30 lines; left out to keep the minimal-stack claim honest. - Sequential prefill on the pure-attention models — and it costs us the first token. Phi-3/Qwen3 prompt tokens still run the full decode path one at a time (cross-turn prefix reuse makes this cheap for multi-turn chat — only the delta prefills). Measured consequence: on a cold ~35-token prompt our TTFT is 291 ms (Phi-3) / 453 ms (Qwen3-4B) against WebLLM's implied ~150 ms — WebLLM has the better first token on short prompts, and we publish that. Qwen3.5's hybrid path prefills in chunks of ≤64 since the 2026-07-29 prefill round (
int4_matmul_batched_dynbatched projections + onegdn_recurdispatch per layer per chunk) and is roughly at parity on TTFT; porting the chunked composition to the attention-only specs is future work and is the top item on BENCH.md's levers list. - One decode engine, two configurations.
engine-core.tsis the singlebuildDecodeEngine;validate.htmlruns it unfused/scalar with blocking per-token readback (deterministic positions, logits access),zero-tvm.htmlruns it fused with URL-flag shader variants and the pipelined readback ring. The per-layer QKV stage and the readback style are the only mode-dependent parts. - Residual buffer ping-pong. WebGPU forbids read+write to the same buffer in one dispatch, so the decode loops swap between
B.residualandB.residual2across theadd_normdispatches. The two swaps per layer cancel out, which is why the per-layer bind groups can be pre-computed once and reused for every token.
MIT. See LICENSE.
This repo ships a CITATION.cff, so GitHub's "Cite this
repository" button renders APA / BibTeX automatically. Each release is archived
to Zenodo — cite the concept DOI
10.5281/zenodo.20838918 for all
versions.
Gunaydin, A. B. (2026). Zero-TVM: Phi-3-mini in the browser on hand-written
WGSL kernels. https://zerotvm.com | https://github.com/abgnydn/zero-tvm
