MLC web-llm path: working in-browser fine-tune blocked by state-leakage bug — parked, documented

[IdkwhatImD0ing]

TL;DR

Parallel investigation to #6. The MLC path produces a smaller, WebGPU-native bundle than ONNX (2.5 GB vs 3.2 GB) thanks to PR #3485's PLE table packing, and all three Gemma 4 variants produce coherent prompt-1 output in-browser after our patches. It is not the shipping path right now because of a confirmed web-llm state-leakage bug that makes multi-call usage unsafe without an engine-reload workaround. This issue documents where the MLC path stands so we can pick it up later (if ONNX hits a wall, or once upstream fixes land).

Goal

Ship the fine-tune to the browser via @mlc-ai/web-llm + WebGPU, replacing transformers.js + ONNX in client/lib/gemma/local-runtime.ts. Concretely:

1. WAVE's various task calls (check-in chat, chunk generation, reflection, insights) all run cleanly with no cross-task state contamination.
2. Cold-load time competitive with current ONNX path (~30s for 2.5 GB into OPFS).
3. Tok/s on Apple WebGPU comparable to current ~15 tok/s via transformers.js, ideally higher (MLC's WebGPU kernels are typically 2-3× faster on modern Apple silicon).

Why we want this

• Smaller bundle: 2.5 GB (q4f16_1 with Gemma4SplitScaledEmbedding packing the PLE tables) vs 3.2 GB ONNX (which can only int4-pack regular Gather, not the PLE Gather via MatMulNBitsQuantizer). PR #3485's quantization is Gemma-4-aware.
• Native WebGPU: MLC compiles to WebGPU shaders directly via TVM. Tok/s on Apple M-series and modern dGPUs should beat transformers.js + onnxruntime-web by a healthy margin.
• Better KV cache memory: TVM's paged KV cache is more memory-efficient than ORT's contiguous approach. Helps on lower-RAM iPhones / Macs running multiple models concurrently (we have Whisper + Gemma + Kokoro in parallel — see client/README.md).
• Future-proof: MLC is actively adding Gemma family support; once PR #3485 lands properly, we benefit from upstream maintenance.

Background — what's already done

Build environment (paid for)

Per docs/postmortems/mlc-build.md, we source-built the full toolchain on Mac at /private/tmp/mlc-workspace/:

• relax-base at commit ac9cf7aac (relax PR #346 — Gemma 4 E2B prerequisites)
• mlc-llm-base at commit fa7cf711 (mlc-llm PR #3485 — Gemma 4 E2B text-only support)
• Both with built dylibs in */build/
• mlc_wasm_runtime.bc at mlc-llm-base/web/dist/wasm/
• Editable installs wired into models/.venv via .pth files. Invoke as python -m mlc_llm <subcommand>.

Setting up this build on Windows is non-trivial (LLVM, emscripten, ~half-day plus Windows-specific gotchas). If we resume MLC work on Windows, expect to redo this from scratch unless we can transfer the built .bc and .dylib artifacts.

Pipeline (works end-to-end)

Per docs/postmortems/mlc-finetune.md:

# 1. PEFT re-merge (same as ONNX path — fixes broken unsloth merge)
uv run --project models python models/merge_lora_peft.py \
  --base unsloth/gemma-4-E2B-it \
  --adapter Maelstrome/lora-wave-session-r32 \
  --out-dir models/runs/merge-peft \
  --device cpu --dtype bfloat16

# 2. Convert weights — peak RAM ~4.6 GB (streaming), output 2.5 GB
uv run --project models python -m mlc_llm convert_weight \
  models/runs/merge-peft \
  --quantization q4f16_1 --device metal:0 \
  --output models/runs/mlc-export-v2

# 3. Generate chat config — then patch the conv_template (see below)
uv run --project models python -m mlc_llm gen_config \
  models/runs/merge-peft \
  --quantization q4f16_1 \
  --conv-template gemma3_instruction \
  --output models/runs/mlc-export-v2

# 4. CRITICAL: patch mlc-chat-config.json conv_template to use <|turn> tokens
# (see "Patches required" section below)

# 5. Compile WebGPU bundle (~13 min on Apple Silicon)
MLC_LLM_SOURCE_DIR=/private/tmp/mlc-workspace/mlc-llm-base \
  uv run --project models python -m mlc_llm compile \
  models/runs/mlc-export-v2/mlc-chat-config.json \
  --device webgpu \
  --output models/runs/mlc-export-v2/wave-r32-q4f16_1-webgpu.wasm

MLC_LLM_SOURCE_DIR (not MLC_LLM_HOME) is checked by auto_target.py:241 — without it, compile fails with Cannot find library: mlc_wasm_runtime.bc.

Patches required to ship (already applied locally)

Patch 1: mlc-chat-config.json conv_template for Gemma 4's actual tokens.

PR #3485 doesn't ship a gemma4_instruction template. gen_config requires picking one; gemma3_instruction is the closest existing but its strings don't match Gemma 4's tokenizer. Verified directly against google/gemma-4-E2B-it's official tokenizer:

Token ID	Gemma 3 (per gemma3_instruction)	Gemma 4 (actual, all variants)
105	<start_of_turn>	<|turn>
106	<end_of_turn>	<turn|>

So replace the generated conv_template block with:

"conv_template": {
  "name": "gemma4_turn",
  "system_template": "{system_message}",
  "system_message": "",
  "system_prefix_token_ids": [2],
  "add_role_after_system_message": true,
  "roles": {
    "user": "<|turn>user",
    "assistant": "<|turn>model"
  },
  "role_templates": {
    "user": "{user_message}",
    "assistant": "{assistant_message}",
    "tool": "{tool_message}"
  },
  "messages": [],
  "seps": ["<turn|>\n"],
  "role_content_sep": "\n",
  "role_empty_sep": "\n",
  "stop_str": ["<turn|>"],
  "stop_token_ids": [1, 106],
  "function_string": "",
  "use_function_calling": false
}

No recompile needed — mlc-chat-config.json is read at runtime by web-llm.

Patch 2 (browser only): Next.js rewrite for /resolve/main/.

web-llm constructs HF-style fetch URLs ${model}/resolve/main/${file}. Add to client/next.config.ts:

async rewrites() {
  return [
    { source: "/mlc-export/resolve/main/:path*",  destination: "/mlc-export/:path*" },
    { source: "/mlc-base-export/resolve/main/:path*",  destination: "/mlc-base-export/:path*" },
    { source: "/mlc-google-it-export/resolve/main/:path*",  destination: "/mlc-google-it-export/:path*" },
  ];
}

Patch 3 (dev only): hard-copy weights into client/public/, don't symlink.

Turbopack panics on symlinks that escape client/ (FileSystemPath("").join("../models/runs/...") leaves the filesystem root). Use rsync -a --exclude='resolve' <model-dir>/ client/public/<model-name>/.

Verified working (prompt-1 only, see blockers below)

The 3-way compare page loads all three Gemma 4 variants through PR #3485 + our patches:

Model	Prompt-1 output
Our fine-tune (PEFT-merged LoRA)	"1, 2, 3, 4, 5" / "The capital of France is Paris." / coherent haikus / coherent anxiety advice
unsloth/gemma-4-E2B-it	Same — coherent factual + creative + therapeutic
google/gemma-4-E2B-it	Same — official upstream weights work

All three at ~45 tok/s on Apple M-series WebGPU.

Blocker — web-llm state leakage

This is the only thing keeping MLC out of production. Documented exhaustively in docs/postmortems/mlc-finetune.md#6.

Symptom: sequential engine.chat.completions.create({messages: [{role: "user", content: prompt}]}) calls leak KV-cache state between calls. The first call's output is coherent; subsequent calls' outputs are contaminated by the first call's context.

Prompt order (with prompt #N as the FIRST call after engine load)	Output
#1 "Count from 1 to 5"	1, 2, 3, 4, 5 ✓
#2 "I'm feeling anxious..." (after Count)	0 tokens emitted ✗
#3 "Capital of France?" (after Anxious)	"Please provide the sentence or question you would like me to answer." ✗
#4 "Write a haiku..." (after Capital)	degenerate wave deep wave deep loop ✗

When the same prompts are run with the order shuffled, whichever prompt runs first is always coherent, and the rest degrade.

Diagnosis: engine.resetChat() IS called by web-llm's prefill path when conversations don't match (verified by reading engine.js). The implementation calls pipeline.resetKVCache() which calls vm.builtin.kv_state_clear + vm.builtin.kv_state_add_sequence. On paper this is a full reset. Empirically it isn't — confirmed by inserting explicit engine.resetChat() between calls (no behavioral change). The actual bug appears to be in TVM's paged KV cache C++ runtime function.

Confirmed by existing upstream issues (don't file another):

• web-llm #215 — resetChat() causing Mistral model to output scrambled responses. Same symptom, different model. Open since 2024, no fix.
• web-llm #735 — Caching Problems. Documents the workaround: pass the full conversation array each call instead of relying on reset.

Workaround that does fix it: full engine reload between each chat.completions.create() call (engine.unload() + CreateMLCEngine(...)). ~3-5s warm-cache overhead. Verified — all four prompts produce coherent output when wrapped in this pattern. See the diagnostic loop in client/app/training/mlc-test/compare-all/compare-all-client.tsx.

What needs to be done

Two parallel paths. Either or both, depending on timeline:

Short-term: production wiring with the engine-reload workaround

Owner: us (when we revisit). Estimated effort: 1-2 days.

1. Refactor client/lib/gemma/local-runtime.ts to swap @huggingface/transformers + @browser-ai/transformers-js for @mlc-ai/web-llm. Two integration points:
   • The simple generateChatText path (chunk gen, insights, reflection, voice-test) — direct chat.completions.create swap.
   • The check-in path uses Vercel AI SDK + tool calling for the endConversation structured tool. No web-llm equivalent exists — would need to either prompt the model to emit JSON and parse client-side, or wrap web-llm in a custom AI SDK provider. Existing parseCheckInPayload + inferEndConversationFromHistory fallback at client/lib/gemma/local-runtime.ts:545+ gives us a JSON-only path.
2. Implement the hybrid engine-lifecycle pattern:
   • One shared engine instance globally.
   • Within a single task type's session (e.g. a multi-turn check-in chat), pass full conversation history per call → web-llm's "multiround chatting" path reuses KV cache cleanly per engine.js:12982.
   • When switching task TYPE (chunk → reflection → insights, etc.), call engine.unload() + CreateMLCEngine() to fully reset. ~3-5s cost on task switch only.
3. Push models/runs/mlc-export-v2/ (2.5 GB) to a new HF repo (e.g. Maelstrome/wave-r32-mlc) so production fetches don't depend on a local dev server.
4. Replace client/public/mlc-export/ hard copy with the HF URL in the production AppConfig.

Long-term: upstream contributions to remove the patches

Owner: open source. Best case: someone files / lands these. Worst case: we maintain the patch stack.

1. mlc-ai/mlc-llm: register a gemma4_instruction conv_template in python/mlc_llm/conversation_template/gemma.py using <|turn>user / <|turn>model / <turn|> / [1, 106]. Either make it the default for model_type=gemma4 or document --conv-template gemma4_instruction. This would be a small, clean PR contribution worth making if we have a couple hours. Reference the same template patch we already use.
2. mlc-ai/mlc-llm or mlc-ai/relax: root-cause and fix the TVM paged KV cache clear bug (the underlying state leak). Out of scope for us — likely requires C++ runtime debugging.

Things to NOT re-litigate

These are decided. See linked memory notes and postmortems.

1. Don't apply self.scaling = 1.0 / math.sqrt(local_head_dim) patch to gemma4_model.py:430. PR #3485's self.scaling = 1.0 is correct — MLC's attention op applies the sqrt internally. Patching double-applies and kills generation. Verified empirically by reverting our own (wrong) patch.
2. Don't use Maelstrome/lora-wave-session-r32-merged as the source. Broken at the safetensors level. Re-merge with PEFT via models/merge_lora_peft.py.
3. Don't use gemma3_instruction conv_template for Gemma 4. Tokens are different despite same IDs.
4. Don't try MLC PR #3485 with a webgpu-only test. Both WebGPU and Metal builds are affected by the same state-leak bug — Metal isn't a workaround.

Reference material

• docs/postmortems/mlc-build.md — source build of the toolchain
• docs/postmortems/mlc-finetune.md — full investigation, patches applied, state-leakage diagnosis
• docs/postmortems/README.md — postmortem index
• models/runs/mlc-export-v2/ (gitignored) — the working 2.5 GB MLC bundle for our fine-tune
• client/app/training/mlc-test/compare-all/ — 3-way browser comparison page
• Upstream web-llm issue #215 — the resetChat bug we're hitting
• Upstream web-llm issue #735 — workaround documentation

Done criteria (if/when we ship MLC)

• client/lib/gemma/local-runtime.ts uses @mlc-ai/web-llm for all task calls
• Engine-reload-between-task-types pattern implemented and tested with a multi-task user flow (check-in → chunk → reflection)
• No cross-task contamination on any prompt across at least 5 sequential task switches
• Within a single check-in chat, multi-turn conversation correctness (full history passed each call) — verified with a 5-turn smoke test
• Maelstrome/wave-r32-mlc (or equivalent) HF repo created with weight shards + WASM + patched mlc-chat-config.json
• cd client && pnpm exec tsc --noEmit passes
• cd client && pnpm dev → load /training/voice-test or /training/onnx-test/compare → in-browser generation produces coherent text

Decision context

User decision recorded 2026-05-13: skip MLC migration for the current ship, ship via ONNX (issue #6). MLC remains a viable alternative if:

• ONNX hits an unforeseen blocker in production
• We need the tok/s boost from native WebGPU kernels
• Upstream lands the resetChat fix and our patch stack shrinks

[IdkwhatImD0ing]

Update 2026-05-14 — Metal-side Python blocker discovered while trying to isolate the web-llm state-leak root cause

While drafting a Python benchmark plan to settle whether the browser-side degenerate output is a web-llm runtime bug or a mlc-llm PR #3485 bug, we ran the existing wave-r32-q4f16_1-metal.dylib through MLCEngine on Apple Silicon. Engine init completed cleanly in 8 s; first prefill hard-crashes before emitting any tokens:

[02:28:13] /tmp/mlc-workspace/relax-base/src/runtime/metal/metal_module.mm:138
tvm.error.InternalError: Check failed: (state != nullptr) is false:
cannot get state: for function batch_prefill_ragged_kv_kernel
Threadgroup memory size (41568) exceeds the maximum
threadgroup memory allowed (32768)

Apple Silicon caps threadgroup memory at 32 KB per Metal kernel. The PR #3485 compile of Gemma 4 E2B emits batch_prefill_ragged_kv_kernel needing 40.6 KB (27 % over). Metal refuses to create the pipeline state, prefill fails, generation never starts. Pipeline:

• Build artifact: models/runs/mlc-export-v2/wave-r32-q4f16_1-metal.dylib (built 2026-05-13 from relax-base@ac9cf7aac + mlc-llm-base@fa7cf711)
• Repro script: models/runs/bench/mlc_smoke.py
• Full stack trace: MLCEngine.chat.completions.create → ThreadedEngineImpl::RunBackgroundLoop → NewRequestPrefillActionObj::Step → PagedAttentionKVCacheObj::AttentionWithFusedQKV → TIRRaggedPrefillFunc::MHA → MetalModuleNode::GetPipelineState → check fail.

What this changes about the parent issue's diagnosis

1. The web-llm browser path doesn't crash with this error on the same Mac — it produces the degenerate '1. **feel** your **feet** **now**. 1. 1.' output we documented, then continues. So WebGPU's MSL compile path is either silently re-sizing the kernel, picking a different variant, or producing wrong values when running the oversized kernel. The browser is hiding a hard failure as garbage output, which is consistent with the per-layer-input scale and attention-scale concerns Gemini flagged in PR #3485's review.
2. The MLC-vs-ONNX Python benchmark plan cannot run on Mac until either upstream lands a Gemma-4 schedule with smaller threadgroups for Apple Silicon, or we patch the TIR schedule locally (out of scope).
3. The plan is still viable on Windows + CUDA: CUDA shared-memory limits are 96-228 KB depending on SM arch, well above 40.6 KB. When the #6 Windows agent lands, they should be able to run the same mlc_smoke.py against a fresh CUDA .dll compile and tell us definitively whether the degenerate output is a web-llm browser bug or a mlc-llm kernel bug.

Concrete next steps (added to the Done-criteria pile)

• On Windows/CUDA box, recompile wave-r32-q4f16_1-cuda.dll and run models/runs/bench/mlc_smoke.py — answers root-cause question
• If Windows/CUDA also produces degenerate output → escalate as a new mlc-llm issue tagged at PR #3485, cross-link this comment
• If Windows/CUDA produces coherent output → escalate as a new web-llm issue (separate from #215/#735, since this is wrong-content not wrong-state)
• Separately consider filing the Apple Silicon threadgroup-memory issue against mlc-ai/mlc-llm: Gemma 4 E2B q4f16_1 batch_prefill_ragged_kv_kernel exceeds Apple Silicon Metal threadgroup memory budget (41568 > 32768). Low-priority on our side (we won't ship Mac native), but it's a clean blocker report and there's no public duplicate per the GitHub-search agent's findings.

Update to "Things to NOT re-litigate" #4

The original bullet said "Don't try MLC PR #3485 with a webgpu-only test. Both WebGPU and Metal builds are affected by the same state-leak bug." That's still true, but the Metal native (.dylib) build also can't run prefill at all on Apple Silicon due to the threadgroup memory cap — so "fall back to Metal native to test" was never a viable workaround on Mac for an unrelated reason.