Zero TVM Documentation

Quick start

No install. Just open the chat page. Weights load from your browser cache (if you've used WebLLM before) or download fresh from HuggingFace (~2.1 GB).

How it works

The engine has three phases per generated token:

1. Prefill — process each prompt token sequentially, building up the KV cache
2. First decode — the last prefill step produces the first generated token
3. Decode loop — each step takes the previous token as input, runs the full forward pass with the KV cache providing attention context, produces the next token

Each forward pass runs 10 kernel roles through 32 transformer layers — 228 dispatches per token in the default f16-KV path (260 with the ?int8kv=1 flag) — then reads one i32 token ID back from the GPU. For comparison, WebLLM's TVM-generated decode path fires 342 dispatches per token using 85 distinct shaders.

Architecture overview

Four source files make up the engine:

Weight loader

Weights are stored in MLC's ndarray-cache.json format — an index file listing every parameter, which shard binary it lives in, its byte offset, and byte size. The loader reads this index then fetches each referenced shard.

Fetch priority

1. Browser Cache API — if a prior WebLLM session cached the shards, they load instantly
2. HuggingFace — direct HTTPS fetch from huggingface.co/mlc-ai/Phi-3-mini-4k-instruct-q4f16_1-MLC

Parameter naming (MLC format)

MLC uses non-standard parameter names. The actual names in the cache:

Tokenizer

A hand-written BPE tokenizer in TypeScript. No SentencePiece WASM, no HuggingFace tokenizers bundle. Reads tokenizer.json directly.

Key steps

• Pre-tokenization — Metaspace: spaces become ▁, words are split on whitespace
• BPE encoding — merge pairs by rank from the merge table in tokenizer.json
• Special tokens — <|system|>, <|user|>, <|assistant|>, <|end|>
• Chat template — Phi-3 format applied by buildChatPrompt()

Phi-3 chat template

<|system|>
You are a helpful assistant.<|end|>
<|user|>
What is the capital of Australia?<|end|>
<|assistant|>

Stop tokens: 2 (EOS), 32000 (<|end|>), 32007 (<|endoftext|>).

KV cache

Uses a paged KV cache — the same approach as vLLM and PagedAttention. Memory is divided into fixed-size pages (16 slots each). A page table maps logical positions to physical pages.

Each layer has its own GPUBuffer for KV pages. The page table is a simple identity mapping (page i → physical page i) for single-sequence inference.

Decode loop

Each call to decodeToken(tokenId, position) submits one command encoder with the full forward pass:

// Per-token GPU state written via writeBuffer
B.inputIds   ← [tokenId]          // i32
B.posMap     ← [position]         // i32
B.pageIndptr ← [0, nnzPages]      // page range
B.lengthInfo ← [position+1, 0, 0] // seq length

// Forward pass (one command encoder)
embedding(B.residual)             // token → hidden state
rmsNorm(B.hidden1, B.residual)    // initial norm

for L in 0..32:
  // QKV matmul + RoPE + KV-append in ONE dispatch (M4 fusion)
  qkvFused(B.qOut, kvPages[L], B.hidden1)
  attention(B.attnOut, B.qOut, kvPages[L])
  int4Matmul(B.hidden2, B.attnOut)               // O projection
  addNorm(B.hidden2, resIn → B.hidden1, resOut)  // residual + RMSNorm, ping-pong
  // Gate + Up + SiLU + mul + Down in ONE dispatch
  fusedFfn(B.hidden2, B.hidden1)
  addNorm(B.hidden2, resIn → B.hidden1, resOut)  // residual + RMSNorm, ping-pong

int4Matmul(B.logits, B.hidden1)   // LM head
argmax(B.tokenOut, B.logits)      // → next token ID

Ping-pong residual buffers

WebGPU's validation rules forbid binding the same buffer as both read and read_write in the same dispatch. The add_norm shader needs to read the old residual and write the new one.

Solution: two residual buffers that alternate each dispatch.

let resIn  = B.residual   // ping (starts with embedding)
let resOut = B.residual2  // pong (uninitialized)

// Each add_norm:
dispatch(addNorm, [delta, resIn, gamma, hidden1, resOut])
[resIn, resOut] = [resOut, resIn]  // swap — O(1), no GPU copy

WGSL Kernel Roles

All shaders live in src/compiler/shaders/ — 27 .wgsl files implementing 10 distinct kernel roles. The rest are tiled and subgroup variants of the same role, selectable at runtime via URL flags. The compiler compiles them all at startup into GPUComputePipeline objects.

Binding convention: @group(0) always. Binding indices are zero-based and match the order you pass buffers to bg(device, pipeline, [...bufs]).

Phi-3 model constants

export const PHI3 = {
  D:        3072,   // hidden dimension
  HEADS:    32,     // attention heads
  HEAD_DIM: 96,     // D / HEADS
  LAYERS:   32,     // transformer layers
  FFN:      8192,   // FFN intermediate dimension
  VOCAB:    32064,  // vocabulary size
  PAGE_SIZE:16,     // KV cache slots per page
  MAX_PAGES:257,    // max pages (≈ 4096 context)
}

Q4F16 quantization format

MLC's Q4F16_1 format packs 8 int4 values into each uint32. Scales are stored as float16 with group_size=32 — one scale per 32 weights.

// Extract nibble for element i within a u32
let nibble = (packed_u32 >> (i * 4)) & 0xF;

// Dequantize: center around 0, multiply by scale
let value = f16(i32(nibble) - 7) * scale;

Weight shapes in Q4F16 (for Phi-3-mini):

Port to Phi-4-mini or Qwen3

Both are available as MLC Q4F16 packages. The steps to port:

1. Update PHI3 constants in compiler.ts — D, HEADS, HEAD_DIM, LAYERS, FFN, VOCAB
2. Check parameter names — fetch ndarray-cache.json and log all keys. Update weight-loader.ts candidates to match
3. Check for GQA — if KV heads ≠ Q heads (grouped-query attention), the attention shader needs a small change to repeat KV heads
4. Update chat template — each model has its own special tokens and prompt format
5. Update HuggingFace base URL — change PHI3_MODEL_BASE in weight-loader.ts

Local weight serving

Run the download script once to save all shards locally. Subsequent page loads are instant (served from localhost, no network).

node scripts/download-weights.mjs

# Downloads to: public/weights/Phi-3-mini-4k-instruct-q4f16_1-MLC/
# Served at:    /weights/Phi-3-mini-4k-instruct-q4f16_1-MLC/
# Size:         ~2.1 GB

Debugging tips

All output is <unk>

• Weight version mismatch — ndarray-cache.json offsets don't match shard content
• Buffer aliasing — add_norm dispatched with same buffer as both @1 and @4
• Wrong rope binding order — check @0=q_out, @1=k_out, @2=v_out, @3=qkv, @4=posMap

Garbage / repetitive output

• Rope bindings are in the wrong order (this was our bug — garbage like -,unlintzegesenma)
• Wrong uniform values for a shader (K_groups, N, etc.)

WebGPU validation error about aliasing

• Same buffer bound as read_write and read in one dispatch
• Fix: use ping-pong buffers. Never bind B.residual as both @1 and @4 to add_norm

Model not loading (Weight not found)

• Log all available keys: the weight loader prints them to console on load
• MLC names differ from HuggingFace standard names (transformer.h.N.mixer.* not model.layers.N.self_attn.*)

Bugs we fixed (and how)

vs WebLLM

Head-to-head on Phi-3-mini-4k-instruct Q4F16_1, same weights, same browser (Chrome with WebGPU), Apple M2 Pro:

Where the 22% gap comes from

At ~40 tok/s Zero-TVM already wins on dispatch overhead (228 vs 342 submissions per token). The remaining gap is inside the matmul and attention kernels themselves:

• Matmul tiling. WebLLM's fused_dequantize_NT_matmul kernels are TVM-autotuned with aggressive shared-memory tiling and 2D thread blocks. Our int4_matmul_tiled / _tiled8 / _sg variants close some of this, but the default int4_matmul loses ~15% here.
• Attention reduction. TVM's batch_decode_paged_kv_kernel splits the softmax reduction across more workgroups with a follow-up attention_merge_state pass. Our single-workgroup-per-head design is simpler but leaves M2 Pro's 20 SMs partly idle.
• Subgroup ops. Recent WebGPU supports subgroup reductions; we have _sg variants but they aren't default. WebLLM compiles with subgroup ops turned on.

What is not in the gap: dispatch submission overhead, command-encoder build time, or framework marshalling. Those all favor Zero-TVM and are already measured out.

The memory-bandwidth ceiling

Phi-3-mini Q4F16_1 touches ~1.8 GB of weights per decode token. On M2 Pro's 200 GB/s memory bus that's 9 ms/token, or ~111 tok/s theoretical max. Zero-TVM at ~40 tok/s hits 36% of that ceiling; WebLLM at ~51 tok/s hits 46%. Neither engine can exceed the ceiling without changing weight layout or quantization (int8 KV shaves a bit, which is why it's exposed behind a flag).