diff --git a/AI.md b/AI.md new file mode 100644 index 0000000000000000000000000000000000000000..64cea5c0b7ca8c7472dedc5c677d8096d02abf1e --- /dev/null +++ b/AI.md @@ -0,0 +1,940 @@ +# Everything WebGPU — project harness + +## Project Vision +An embeddable local-LLM engine for the browser: WebLLM on WebGPU, plus a priority scheduler, multi-step +decoding and compute-pass batching that together take decode from 9.7 to 25.9 tok/s. `main` is the library — +host-neutral, no UI, usable from a page, a worker or an extension. `demo` keeps the Firefox WebExtension +that was the original vehicle and is now the library's first consumer. See [ARCHIVE.md](ARCHIVE.md). + +Weights arrive by whichever route suits the app: WebLLM's 163 prebuilt models from HuggingFace, any base URL +the developer hosts, or a local folder read off disk with no network connection at any point. + +**Model: `empero-ai/Qwen3.8-2B-Distill`, compiled to MLC in-house — done and running at 16.6-18.1 tok/s.** The original "4B+" goal was set before decode was instrumented; the measurements retired it. Decode is memory-bandwidth-bound, so time per token scales with weight bytes: a 4B at `q4f16_1` is ~2.25 GB and projects to 7-9 tok/s with only one engine fitting in 16 GB, while this 2B is 1.06 GB and was projected at 14-18 tok/s — the measurement landed inside that band. Build notes and the toolchain fixes are in [MLC-COMPILE.md](MLC-COMPILE.md). + +## Workflow & Development Principles +- **Fail Fast**: Validate inputs, model states, and cache availability early. Throw descriptive errors immediately upon invalid conditions. +- **Minimal Surface**: Write only necessary code and tests. Prefer a single, comprehensive integration test over redundant unit tests. +- **Reuse First**: Leverage existing internal APIs, built-in libraries, and ecosystem patterns (e.g., WebLLM, Cache API) before introducing net-new abstractions. +- **Direct Execution**: Output exact code changes or direct answers. Omit preamble, pleasantries, conversational fillers, and unsolicited caveats. + +## Where the work lives + +| | | +| --- | --- | +| [README.md](README.md) | The developer-facing entry point. Asserted by [test/readme.test.mjs](test/readme.test.mjs), so its examples cannot drift from the API. | +| [API.md](API.md) | Every call form, one page. Asserted by [test/api-doc.test.mjs](test/api-doc.test.mjs) — method names, error codes, enum values and the export list are all derived from the source. | +| [ROADMAP.md](ROADMAP.md) | **The only list of open work.** | +| [ARCHIVE.md](ARCHIVE.md) | What was done and *why* — decisions with their reasoning, so they are not re-litigated. | +| [WEBLLM-SURFACE.md](WEBLLM-SURFACE.md) | What WebLLM already does. **Read before adding a capability**, and follow its "Upgrading" runbook on every dependency bump. | +| This file | What is true and measured. Reference, not a plan. | + +Task lists used to live here, in AI2.md and in NATIVE-REUSE-PLAN.md at the same time, with "Track 1" +and "Track 2" meaning different things in each. That is how a reader implements the wrong item, and +it is the same shape of failure — no single place to look before acting — that produced the +duplicated WebLLM helpers ARCHIVE.md records. + +## Completed Tasks + +Through the in-house model compile. Everything after that — the library extraction, the model +sources, the WebLLM de-duplication — is in [ARCHIVE.md](ARCHIVE.md) with its reasoning. + +- [x] Defined core requirements for local WebGPU-based execution in Firefox on macOS. +- [x] Established direct cache-injection architecture for offline local models. +- [x] Configured `manifest.json` (MV2, persistent background page, `wasm-unsafe-eval` CSP, `unlimitedStorage`) and documented the `about:config` flags in the manager page. +- [x] Implemented drag-and-drop model folder ingestion that writes straight into `Cache Storage` under WebLLM's own scopes and keys. +- [x] Built the minimal test chat popup driven by the background WebLLM engine. +- [x] Exposed the `everything-webgpu/v1` message + port API for other Firefox extensions. +- [x] Verified end-to-end on real hardware with `Qwen3.5-0.8B-q4f16_1-MLC` (Firefox 154, macOS, M4). +- [x] Root-caused decode throughput to Firefox's 100 ms WebGPU poll timer (Mozilla bug 1870699), not to anything in this extension. +- [x] Replaced the fail-fast `busy` flag with a priority scheduler: bands, session supersession, opt-in preemption. +- [x] Added a worker-backed engine pool and a `batch` API so independent work runs concurrently. +- [x] Implemented multi-step decoding (vLLM's `--num-scheduler-steps`, default 15) so one GPU sync covers K tokens: 9.7 -> 18.4 tok/s single-stream. +- [x] Instrumented decode with a CPU-encode / GPU-execute probe and root-caused the remaining ~46 ms/token to 664 per-token kernel launches, each in its own WebGPU compute pass — not the poll tick, not command encoding. +- [x] Batched consecutive tvmjs kernel launches into one compute pass (664 passes/token -> ~16): 10.3 -> 25.9 tok/s on an identical greedy generation, byte-identical output. +- [x] Priced the remaining gap against the platform: ~50 GB/s is the most a hand-written WGSL kernel gets here, ~16 GB/s is what the generated dequant-GEMV gets, and the difference is the reduction tail amortised over a 2-iteration loop — not load width (~1.15x) and not dequantisation (~free). +- [x] Compiled `Qwen3.8-2B-Distill` to MLC/WebGPU in-house: vision tower and MTP head stripped, `q4f16_1`, 1.06 GB, 26 shards. Five separate breakages in the published MLC nightlies had to be worked around first — see [MLC-COMPILE.md](MLC-COMPILE.md) and `tools/`. +- [x] Made the engine pool demand-driven: one engine at load, a second only when a second task competes, one engine per task. Measured that a second engine is worth 1.06x on this model, so it buys isolation rather than throughput. +- [x] Surfaced the model-load report in the UI. WebLLM's text was always being passed through, but `#status` was a single ellipsised line, so it was truncated to `Loading model from cache[26/5…`. + +## Consolidated Context +- **Target Platform**: Firefox WebExtension (macOS, requiring WebGPU flags). Test machine: M4 MacBook Air, 16 GB unified memory, ~120 GB/s. +- **Core Stack**: JavaScript, WebGPU, WebLLM (`@mlc-ai/web-llm` 0.2.84, patched at build time), Cache API (for local file injection), Extension Message Passing. +- **Architecture**: A host-neutral engine (`src/engine/`, asserted free of any WebExtension API) behind + per-host adapters (`src/adapters/`). The Firefox extension is one host: background engine host + popup test + UI + extension-to-extension API provider. See [ARCHIVE.md](ARCHIVE.md). +- **Model**: `Qwen3.8-2B-q4f16_1` (1.06 GB), compiled in-house from `empero-ai/Qwen3.8-2B-Distill` — see [MLC-COMPILE.md](MLC-COMPILE.md). `Qwen3.5-0.8B-q4f16_1-MLC` remains the baseline most of the analysis below was measured on. Keep `q4f16_1` — dequantisation measured ~free, so wider formats only add bytes, and bytes are what decode pays for. +- **Scheduling**: one shared GPU, one engine per task, pool grows on demand. A second engine measured 1.06x on this model, so it buys isolation rather than throughput. +- **Build-time patches** ([build/patches.mjs](build/patches.mjs), applied by [build.mjs](build.mjs)): `storage-buffer-limit` (Firefox caps storage buffers per stage at 9, tvmjs asks for 10) and `compute-pass-batching` (one compute pass per kernel launch -> one per flush). Anchors are literal JS matched modulo whitespace, word-bounded, and optionally scoped to an enclosing function; all are verified before anything is rewritten, so a WebLLM bump reports every break at once with the nearest candidate lines. `npm run verify-patches` checks them without rebuilding; `NO_PASS_MERGE=1` skips the second patch for A/B. + +--- + +## Verified + +All numbers on an M4 MacBook Air (16 GB), Firefox 154 release, macOS. + +### The shipping model: `Qwen3.8-2B-q4f16_1` (1.06 GB, 26 shards) + +Compiled in-house; see [MLC-COMPILE.md](MLC-COMPILE.md). + +| | | +| --- | --- | +| Ingest 1.06 GB into Cache Storage | 4.8 s | +| Model load (cache only, zero network) | 51 s | +| Prefill | 48 tok/s short prompt, 100-200 tok/s at length | +| **Decode** | **16.6-18.1 tok/s** | +| Re-prefill cost per history token | 5.27 ms (no cross-turn KV reuse — see Current Tasks) | +| Kernel launches per decoded token | 664 = 639 forward + 25 sampling, across ~16 flushes | +| Second engine, 4-prompt batch | **1.06x** — see "Scheduling" | +| Two tasks on two engines (e2e) | 1.05x, then **1.11x and 1.12x** on two later runs. The ratio held across runs whose absolute times differed by 35%, so it is a real shift rather than noise — but 1.11x is still isolation, not throughput scaling. | + +### The baseline it was built against: `Qwen3.5-0.8B-q4f16_1` (443 MB, 11 shards) + +Everything below this point was measured on the 0.8B. It is kept because it is where the +architecture came from — the ceiling, the multi-step fix and the pass-batching fix were all found here. + +| | | +| --- | --- | +| WebGPU in the MV2 background page | available, `shader-f16` supported | +| Ingest 443 MB into Cache Storage | 2.1 s | +| Model load (cache only, zero network) | 48 s | +| Prefill | 95-98 tok/s | +| Decode, stock single-step | 9.6 tok/s (warm) | +| Decode, multi-step K=15, one pass per kernel | 17.3-18.3 tok/s | +| Decode, multi-step K=15 + batched compute passes | **24.9-28.0 tok/s** (25.9 when first measured; 27.4, 28.0, 24.9 across later e2e runs on the same build — run-to-run spread is ~12%, so treat any single number as ±1.5) | +| Decode, K=32 + batched passes | 26.8 tok/s | +| Kernel launches per decoded token | 664, across ~16 flushes | +| Decode budget at K=16 | 3.8 ms CPU encode + ~34 ms GPU + <6 ms tick | + +Reproduce with `npm run e2e`, or `ENGINE_COUNT=1 DECODE_STEPS=1,2,4,8,16,24 npm run e2e` for the multi-step +curve — see [test/e2e/run.mjs](test/e2e/run.mjs). + +Three things that surfaced from running it for real: + +- **WebGPU works in the background page.** This was the load-bearing assumption behind putting the engine + there, and it holds on release Firefox. +- **Neither the poll tick nor CPU encoding is the ceiling.** Both measured small once instrumented. What was + left split in two: 664 per-kernel compute passes (fixed, ~2.5x) and, underneath, decode running at ~13% of + memory bandwidth (open). See "Where the 46 ms goes" and "Still not 100+ tok/s". +- **Firefox needed a shim.** tvmjs hardcodes a request for 10 storage buffers per shader stage; Firefox's + Metal backend caps `maxStorageBuffersPerShaderStage` at 9, so `detectGPUDevice()` threw before a device was + ever requested. `build.mjs` clamps that request to what the adapter reports, and fails the build loudly if + the patch stops matching after a WebLLM upgrade. Kernels that genuinely need a 10th binding still fail at + pipeline creation — **silently**, as a no-op dispatch that emits garbage rather than an error. Both models + ship four such kernels; three are unreachable by config and the fourth, `batch_prefill_paged_kv_kernel`, is + kept off the live path by `engine-worker.js` calling `resetChat()` before every prefill when the device + reports fewer than 10. `node tools/audit-wasm.mjs ` checks this. + +### The 10 tok/s ceiling + +Decode is capped by Firefox, not by this extension or by the GPU. [`WebGPUParent`](https://searchfox.org/firefox-main/source/dom/webgpu/ipc/WebGPUParent.cpp) +detects GPU completion by polling on a fixed timer: + +```cpp +const uint64_t POLL_TIME_MS = 100; +mTimer.Start(base::TimeDelta::FromMilliseconds(POLL_TIME_MS), this, &WebGPUParent::MaintainDevices); +``` + +So every `onSubmittedWorkDone()` / `mapAsync()` resolves only on a 100 ms tick. WebLLM needs exactly one +GPU->CPU sync per token (it reads back the 4-byte sampled token id before it can compute the next one), and a +serial sync loop always lands just after a tick and waits the full period. One token per tick = 9.6 tok/s. + +`npm run bench` (no extension, no model, plain page) measures this directly: + +| measurement | result | meaning | +| --- | --- | --- | +| `idleSyncMs` | 104 ms | awaiting an **empty** queue with nothing submitted still costs a full tick | +| `dispatch256Ms` | 104 ms | 256 dispatches **sharing one compute pass and one bind group** cost the same as zero | +| `pipelined10SubmitsMs` | 104 ms | 10 submits, one await — it is latency, not throughput | +| `syncAfterRandomDelay` | mean 52 ms, min 3, max 106 | desynchronize from the tick and you get the expected uniform 0-100 ms | +| `n2048` | `1pass=104ms 2048passes=309ms` | a compute *pass* costs ~100 µs | +| `dispatch65536InOnePassMs` | 207 ms | a *dispatch* inside a pass costs ~2 µs — measured past the tick, not bounded by it | + +Read the `dispatch256Ms` row carefully — reusing one pass and one bind group is **not** what the runtime does, +and reading it as "compute is free" is what produced the wrong conclusion corrected below. + +WebLLM's own `enable_latency_breakdown` agrees: `totalTime=101.2ms`, of which `sampleTime=101.1ms` — the step +containing the sync is the entire token budget. Identical numbers in a hidden background page, a visible tab, +and an ordinary web page, with and without `MOZ_DISABLE_GPU_PROCESS=1`, so it is not extension-specific and +not GPU-process IPC. + +> **Correction.** This section used to claim actual compute was "~7 ms/token (≈137 tok/s if the sync were +> free)". That was wrong, and it was wrong in a way that mattered: the 7 ms came from `dispatch256Ms - +> dispatch1Ms`, a micro-benchmark of trivial kernels sharing one compute pass and one bind group — nothing +> like a real forward pass. Removing the per-token sync (see "Multi-step decoding") took decode to 18.4 tok/s, +> not 137, because real per-token GPU work is ~46 ms. See "Where the 46 ms goes". + +This is [Mozilla bug 1870699, "Don't poll WebGPU from a timer"](https://bugzilla.mozilla.org/show_bug.cgi?id=1870699), +still open on trunk. Event-driven completion in Firefox would fix it outright. Short of that, the lever is to +stop needing a sync per token — see "Multi-step decoding" below, which does exactly that and lifts a single +stream well past 10 tok/s. + +**Practical consequence:** the tick sets a *floor* of ~10 tok/s on serial decode, and multi-step lifts a +single stream off it — but only to ~18 tok/s, where a second, larger cost takes over. Prefill is unaffected +(one sync per forward pass regardless of token count), so prompt processing, embeddings and short completions +were never the problem. + +### Multi-step decoding + +The ceiling above is a *sync* budget, so the fix is to spend fewer syncs: run K forward steps per sync. This +is vLLM's [`--num-scheduler-steps`](https://blog.vllm.ai/2024/09/05/perf-update.html) (v0.6.0, +28% on +Llama-70B / 4xH100), and it lands here without recompiling any model, because WebLLM's sampling path is +already entirely on the GPU. `softmax_with_temperature` -> `argsort_probs` -> `sample_with_top_p` returns an +int32[1] **device** tensor, and `Tensor.copyFrom(Tensor)` is a device-to-device copy — so the sampled id can +be fed straight back into `embed` without ever becoming a JS number: + +``` +embed -> decode -> penalties -> softmax -> argsort -> sample ─┐ x K, no sync + ^ │ + └──────────────── device tensor, never read back ───────────┘ + then: one device.sync() +``` + +Each step stages its id into its own CPU tensor; tvmjs queues those readbacks into `pendingGPUToCPUCopy` and +only awaits them in `sync()`, so K readbacks still cost one tick. `steps` defaults to **15**, vLLM's documented +cap. [src/background/multistep.js](src/background/multistep.js) is the whole implementation; it lives in the +engine worker because that is where the decode loop is. + +**Measured, single engine, Qwen3.5-0.8B on an M4 Air.** `DECODE_STEPS=... npm run e2e` sweeps K live on one +loaded model and reports the probe built into the burst — the K-step loop contains no `await`, so timing +either side of its single sync partitions the budget exactly: + +| K | tok/s | cpu-encode | gpu + tick | kernels/tok | +| --- | --- | --- | --- | --- | +| 1 (stock) | 9.7 | — | — | — | +| 2 | 11.5 | 4.3 ms | 82.1 ms | 664 | +| 4 | 16.0 | 4.8 ms | 57.6 ms | 664 | +| 8 | 16.8 | 4.5 ms | 54.8 ms | 664 | +| 16 | 17.8 | 4.2 ms | 52.1 ms | 664 | +| 24 | **18.4** | 4.2 ms | 50.1 ms | 664 | + +The curve is smooth and monotone, saturating near 18-19 tok/s: **+90% over stock at the shipped default of +15**. There is no sawtooth — an earlier version of this section predicted one from tick quantization, which +only appears when per-step compute is small next to the 100 ms tick. It is not: it is ~46 ms. Past K≈16 the +tick is already amortized to under 6 ms/token and there is nothing left to win. + +### Where the 46 ms goes + +Subtracting the amortized tick from the `gpu + tick` column leaves ~46 ms/token of real GPU time, against a +memory-bandwidth floor of ~3.5 ms (420 MB of q4f16 weights read once per token at ~120 GB/s). So ~93% of +decode is overhead, and the probes say which: + +| candidate | measured | verdict | +| --- | --- | --- | +| the 100 ms poll tick | <6 ms/token at K≥16 | amortized away by multi-step | +| CPU command encoding, `createBindGroup`, IPC | 4.2 ms/token, flat in K | **not** the bottleneck | +| GPU execution | ~46 ms/token over **664 kernel launches** | this is it, at ~70 µs/kernel | + +664 kernels per token is ~24 per layer — an unfused graph — and tvmjs opens a **separate compute pass for +every one of them** (`submitShader`: `beginComputePass` → `setPipeline` → `createBindGroup` → `dispatch` → +`end`). WebGPU barriers between passes, and on Metal each pass is its own command encoder. `npm run bench` +isolates the cost: + +| measurement | result | +| --- | --- | +| 2048 dispatches, **one** compute pass | 104 ms (i.e. free — it is all tick) | +| 2048 dispatches, **2048** compute passes | 309 ms → ~100 µs per pass | +| CPU cost of encoding one tvm-style kernel | 3.9 µs | + +So the pass is ~25x its own encode cost, and it is charged 664 times per token — while only **16 flushes per +token** actually happen, i.e. ~41 consecutive launches share an encoder and each open their own pass for no +reason. Dispatches *within* one WebGPU compute pass are already ordered with implicit barriers (a compute +pass has a usage scope per dispatch, so implementations must synchronize between them), so those passes were +not buying correctness. + +### Batching the compute passes + +The `compute-pass-batching` patch in [build/patches.mjs](build/patches.mjs) rewrites tvmjs to open a pass +lazily and close it in `flushCommands()` — already the one chokepoint every operation that cannot run +mid-pass routes through. That turns 664 passes/token into ~16. Three edits, each refusing to apply if its +anchor stops matching after a WebLLM upgrade, same as the storage-buffer shim next to it. The `compute.end();` +anchor is scoped to the function the pass is opened in, so an unrelated compute pass elsewhere in tvmjs is not +mistaken for ambiguity. Build with `NO_PASS_MERGE=1` to skip it and A/B on one machine. + +Measured on the same 127-token greedy generation (`temperature: 0`, single engine, K=15): + +| | stock WebLLM | + multi-step | + pass batching | +| --- | --- | --- | --- | +| decode | 9.7 tok/s | 10.3 tok/s | **25.9 tok/s** | +| `gpu + tick` | — | 85.2 ms/tok | **33.8 ms/tok** | +| `cpu-encode` | — | 5.1 ms/tok | 4.4 ms/tok | +| kernels/token | 664 | 664 | 664 | +| compute passes/token | 664 | 664 | **~16** | + +Output is **byte-identical** across the A/B (490 bytes, greedy) — the check that matters, since the whole +change rests on WebGPU synchronizing dispatches inside a pass. + +One caveat on the numbers: the unbatched build is *noisy* run to run (K=16 measured anywhere from 10.7 to +18.6 tok/s, `gpu + tick` from 49.6 to 92.1 ms), while the batched build is tight (24.6-26.8 tok/s, 33.4-37.3 +ms across five runs). The table above is a same-session pair, and its baseline happened to land on the slow +side. Against the *best* baseline seen the gain is ~1.4x rather than 2.5x; against the median, ~1.8x. The +direction and the mechanism are not in doubt, but quote a range, not the 2.5x. + +The K curve keeps the same shape, shifted up, and still has no sawtooth: 2 → 19.3, 8 → 23.8, 16 → 25.6, +32 → 26.8 tok/s. `DEFAULT_DECODE_STEPS` stays at 15 (25.6 tok/s); K=32 buys another 4.7% for double the +transient buffers and more lookahead discarded at every stop token. K=1 is unchanged at 9.7 — with GPU time +now ~34 ms it fits inside a tick, so single-step is purely tick-bound and only multi-step can help it. + +### Still not 100+ tok/s, and it is not the launch count + +~34 ms/token of GPU remains. It is tempting to blame the 664 launches, and an earlier version of this file +did. That is wrong, and the bench says so once it is pushed past the tick: + +| dispatches in **one** pass | wall | per dispatch | +| --- | --- | --- | +| 2 048 | 104 ms | ≤50.8 µs — but this *is* the tick, so it only bounds the number | +| 16 384 | 102 ms | ≤6.2 µs — still one tick | +| 65 536 | 207 ms | ~2.1 µs marginal | + +A dispatch inside a pass costs **~2 µs**, so all 664 of them cost ~1.3 ms, and the ~16 remaining passes add +~1.6 ms. Neither is the 34 ms. Note how the first row alone would have supported the wrong conclusion: any +measurement that lands on 104 ms is the poll grid talking, not the GPU. + +What is left is the kernels' own execution. Splitting the launches at the forward/sample boundary gives +**639 forward + 25 sampling** — the 248k-vocab `argsort_probs` is only 4% of the budget, so the sampling tail +is not it either. Subtracting the amortized tick leaves ~27 ms of forward execution to move 420 MB of q4f16 +weights, i.e. **~16 GB/s against the M4's ~120 GB/s — about 13% of peak**. + +That is the signature of batch-size-1 decode: a matrix-*vector* product where every weight is read once, used +for a single multiply-accumulate, and thrown away, with 4-bit dequantization on top. Intermediates are not +the problem (hidden size 1024 means a ~2 KB activation, negligible next to the weights). So the lever is +*faster* kernels, not fewer: better dequant-GEMV, or a quantization whose kernels reach a higher fraction of +peak. See "Does a different model help?". + +### Why not llama.cpp/Ollama-class, and what would close it + +Decode moves 420 MB of weights per token, so every question about throughput is a question about achieved +memory bandwidth. `npm run bench` measures the two ceilings above us — a hand-written WGSL kernel that does +nothing but stream a 512 MB buffer, swept over workgroup counts so an untuned kernel is not mistaken for a +platform limit: + +| layer | achieved | of hardware peak | +| --- | --- | --- | +| M4 (128-bit LPDDR5X) theoretical | ~120 GB/s | 100% | +| Best hand-written WebGPU streaming read, Firefox | **~50 GB/s** | ~43% | +| MLC/TVM dequant-GEMV during decode | **~16 GB/s** | ~13% | + +Two independent gaps, with different owners: + +- **120 → 50 is the platform.** Flat at 50-53 GB/s across 256/1024/4096/16384 workgroups, so it is not an + occupancy or parallelism problem in the probe. Nothing in this repo moves it. +- **50 → 16 is the kernels.** This is the one worth attacking: ~3x, and it is entirely inside the generated + dequant-GEMV. + +#### Where the 50 -> 16 goes, measured + +The generated kernel is in the model's `.wasm` as WGSL, so it can be read directly. The hot one +(`fused_dequantize2_fused_NT_matmul1_silu1_multiply1_kernel`) is `@workgroup_size(64)`, each thread runs a +**2-iteration** loop over packed weights, and the workgroup then reduces through `var red_buf0` +with barriers. + +`npm run bench` rebuilds that kernel one property at a time against the same 512 MB buffer, so each step's +cost is isolated: + +| kernel property | GB/s | cost | +| --- | --- | --- | +| `vec4` loads, 256-wide, long loop — the ceiling | ~50 | — | +| scalar `u32` loads (what packed 4-bit weights force) | ~41-46 | ~1.15x | +| + unpack 8 nibbles, centre, scale | ~41 | ~1.0x — **dequant arithmetic is nearly free** | +| + 64-lane barrier reduction, 32 iters/thread | 34.0 | 1.2x | +| + 8 iters/thread | 27.0 | 1.5x | +| **+ 2 iters/thread — what MLC actually generated** | **18.6** | **2.2x** | +| decode in the real engine (adds scale traffic, writes, real grid) | ~16 | 2.6x | + +So neither load width nor dequantisation is the problem. **The cost is the reduction tail amortised over +almost no work**: each workgroup reads 128 `u32` — 512 bytes — and then spends six `workgroupBarrier()`s +doing a 64-lane tree reduction. Raising work per thread from 2 to 32 iterations recovers **1.83x** on its own +(18.6 -> 34.0 GB/s) with everything else held fixed. + +That is a *schedule* parameter, chosen by TVM's dlight GEMV rule at model-compile time and baked into the +`.wasm`. It is not something the runtime can change: the launch grid is emitted by the host side of the same +module, so substituting a shader would desynchronise it from its dispatch shape. The fix is to recompile the +model lib — same weights, retuned GEMV tiling — and it is also exactly what subgroups would make moot, since +`simd_sum` reduces 32 lanes without a barrier at all. + +Caveat: these are faithful *models* of the kernel, not the kernel itself; the real one also streams the scale +array and writes results, which is roughly where the last 18.6 -> 16 goes. + +**What is structurally different in llama.cpp's pipeline** (reasoned from the two codebases, *not* measured +here — Ollama is not installed on the test machine): + +1. **Hand-written Metal vs WGSL→naga→MSL.** WebGPU mandates bounds-checked buffer access; naga emits clamps + on dynamic indexing, which lands in a GEMV's innermost loop. +2. **No subgroup ops.** llama.cpp's Metal GEMV leans on `simd_sum`/`simd_shuffle` for cross-lane reduction + and `simdgroup_matrix` for matmul. WGSL subgroups are not exposed in Firefox, so TVM must reduce through + workgroup shared memory and barriers — the expensive path for a reduction-shaped kernel. +3. **Quantization co-designed with its kernel.** `Q4_K` and friends are laid out so a block dequantizes in + registers from vectorized loads; `q4f16_1` goes through a generated TVM dequant instead. + +**Is "80% of Ollama" reachable here?** Probably not on Firefox WebGPU today, and the arithmetic says why +rather than the vibes. Even if the dequant-GEMV became as efficient as a *pure streaming read* — which it +cannot be, since it also has to dequantize and reduce — 420 MB at 50 GB/s is 8.4 ms/token, about **110-120 +tok/s ceiling** for this model on this platform. Any Ollama number above ~140 tok/s therefore puts 80% of it +out of reach no matter how good the kernels get, because the 120 → 50 gap is not ours to close. + +So the honest targets are: **~70 tok/s is plausible** (halve the 50 → 16 kernel gap), ~110 tok/s is the +absolute platform ceiling, and matching a fast native runtime is not on the table without subgroups landing +in Firefox or a move off WebGPU. Measure the actual Ollama figure on the same machine before committing to a +percentage target — it is the one number in this section that is assumed rather than measured. + +> Also worth knowing: the bench loses the WebGPU device outright when it creates 2048 compute passes in one +> encoder (`deviceLostDuringBench = yes, during pass-sweep`), and a lost device does not throw — later calls +> silently no-op and report impossible numbers like 1374 GB/s. The unpatched runtime issued 664 passes per +> token, which is uncomfortably close to that. The dispatch cap in `patchComputePassBatching()` exists for +> this reason. + +### Does a different model help? + +Decode cost per token is roughly `weight bytes / effective bandwidth + ~1.3 ms of launches + tick/K`. +The measured effective bandwidth is ~16 GB/s of the M4's ~120 GB/s, so the first term dominates and scales +with parameter count and quantization — not with how the layers are arranged. Extrapolating from the measured +26 tok/s at 420 MB: + +| change | weight bytes | projected | why | +| --- | --- | --- | --- | +| another 0.8B architecture, same quant | ~450 MB | ~26 tok/s | same traffic per token; layer layout is not the variable | +| **Qwen3.8-2B-Distill** (shipped) | 1.06 GB | **16.6-18.1 tok/s, measured** | projected 14-18; the projection held | +| Llama-3.2-3B | ~1.7 GB | ~9 tok/s | prebuilt MLC folder exists, so no toolchain — the cheap way to sanity-check these projections | +| a 4B at q4f16 | ~2.25 GB | ~7-9 tok/s | 5x the bytes, and only `engineCount=1` fits in 16 GB. Even at *100%* of peak it is ~19 ms/token = 53 tok/s | +| a ~0.3B at q4f16 | ~160 MB | ~60-70 tok/s | helps, but buys less than fixing the efficiency | +| **same weights, different quantization** | varies | see below | resolved: it changes byte count and nothing else | + +Every row except the 2B is still a projection from the measured 16.8 GB/s. The 2B row is now a measurement, +and it landed inside its projected band — which is the only evidence available that this model of decode cost +predicts anything. + +So: switching models does not get to 100+ tok/s, and going bigger costs throughput roughly in proportion to +the extra bytes. This table is why the "4B+" goal was retired in favour of a 2B — see "Project Vision". +100 tok/s means ~10 ms/token, which at today's 16 GB/s buys only ~160 MB of weights, but at full bandwidth +would buy ~1.2 GB. **Closing the efficiency gap is worth ~7x more than shrinking the model.** + +**Which quantization to recompile with: keep `q4f16_1`.** That experiment is settled by the kernel breakdown +above — unpacking eight nibbles, centring and scaling them measured ~free (41 GB/s with the dequant +arithmetic vs 41-46 GB/s without). Since dequantisation costs nothing and bytes-per-token is what decides +throughput, a wider format is pure loss: + +| format | bytes per 32 weights | vs q4f16_1 | verdict | +| --- | --- | --- | --- | +| `q4f16_1` (current) | 16 + 2 scale = 18 | 1.00x | keep | +| `q3f16_1` | 12 + 2 = 14 | 0.78x -> ~1.28x faster | the only quantization that buys speed, and it costs accuracy | +| `q4f32_1` | 16 + 4 = 20 | 1.11x slower | f32 scales and activations, and f16 ALU is faster on Apple | +| `q0f16` | 64 + 0 = 64 | 3.6x slower | skips a dequant that was already free | +| `q4f16_awq` | 18 | 1.00x | same speed, better accuracy — worth it for quality, not for speed | + +So the ~1.28x from `q3f16_1` is both smaller than the 1.83x available from the schedule *and* the only one of +the two that trades away accuracy. Fix the schedule first; treat quantization as a quality knob. + +One structural note specific to these weights: `vocab_size` is 248320 with `tie_word_embeddings`, so the +output projection alone is 1024 x 248320 = 254 M parameters, about **30% of every token's memory traffic**. +That is unusually large (most models are 32k-128k) and no quantization choice changes its share. + +Set `decodeSteps: 1` to turn multi-step off; `NO_PASS_MERGE=1 npm run build` to turn pass batching off. + +## Build and install + +```sh +npm install +npm run build # bundles @mlc-ai/web-llm into vendor/web-llm.js +npm test # integration test over the cache-injection contract +npm run e2e # real Firefox + real model + real GPU (needs a model folder) +npm run bench # WebGPU submit/sync latency only - no extension, no model +npm run package # -> everything-webgpu.xpi +``` + +Env vars that matter when measuring rather than just running: + +| Var | Applies to | Effect | +| --- | --- | --- | +| `NO_PASS_MERGE=1` | `npm run build` | Skip compute-pass batching, for an A/B on one machine. | +| `ENGINE_COUNT=n` | `npm run e2e` | Force the pool size; use `1` to measure a single stream. | +| `DECODE_STEPS=a,b,c` | `npm run e2e` | Sweep multi-step widths live on one loaded model. | +| `SKIP_BENCH=1` | `npm run e2e` | Drop the two ~40 s `gpuBench` passes when comparing builds. | +| `E2E_VERBOSE=1` | `npm run e2e` | Surface web-ext/Firefox output; without it a failed launch and a hung extension look the same. | +| `MODEL_DIR=…` | `npm run e2e` | Model folder (defaults to `~/Downloads/Qwen3.5-0.8B-q4f16_1-MLC`). | + +Load it with `about:debugging` → This Firefox → Load Temporary Add-on → pick `manifest.json`. + +Before a model can load, set these in `about:config` and restart Firefox: + +| Pref | Value | Why | +| --- | --- | --- | +| `dom.webgpu.enabled` | `true` | Exposes `navigator.gpu`. | +| `gfx.webgpu.ignore-blocklist` | `true` | Only if your Mac's GPU is blocklisted. | +| `dom.webgpu.service-workers.enabled` | `true` | Harmless; needed on builds that gate non-visible contexts. | + +The manager page shows live WebGPU status, so you can tell a flag problem from a model problem. + +## Adding a model + +Three routes, and `load()` resolves across all of them. Two of them are one call: + +```js +// 1. prebuilt — one of WebLLM's 163 HuggingFace models. Nothing to register. +await engine.load("Llama-3.2-1B-Instruct-q4f16_1-MLC"); + +// 2. remote — any base URL you host: an HF repo, a CDN, a path on your origin, localhost. +await engine.registerModel({ + modelId: "Qwen3.8-2B-q4f16_1-MLC", + model: "/models/Qwen3.8-2B-q4f16_1-MLC/", + modelLib: "/models/Qwen3.8-2B-q4f16_1-MLC/Qwen3.8-2B-q4f16_1-webgpu.wasm", +}); + +// 3. local — read off disk. No network connection at any point, ever. +await engine.registerModel({ modelId: "Qwen3.8-2B-q4f16_1-MLC", files }); +``` + +`files` is `{ path, file }[]`; `filesFromDataTransfer` and `filesFromInput` build it from a drop event or a +directory picker. `listAvailableModels()` enumerates all three; `{ prebuilt: false }` on the engine gives a +build that can never fetch a model. + +In the `demo` extension this is the **Models…** page: drop a compiled MLC folder on it. + +A local folder must contain: + +- `mlc-chat-config.json` +- `tensor-cache.json` (or a legacy `ndarray-cache.json`) +- every `params_shard_*.bin` listed in that manifest +- `tokenizer.json` (or `tokenizer.model`) +- exactly one `*-webgpu.wasm` model library + +Grab both halves from Hugging Face — the weights from `mlc-ai/-MLC`, the matching library from +`mlc-ai/binary-mlc-llm-libs` — or compile your own with `mlc_llm convert_weights` + `gen_config` + `compile`. + +Local registration validates the whole folder **before** writing anything, then copies each file into Cache +Storage. A missing shard fails in milliseconds rather than after 2 GB of copying. A remote URL is not +validated at all — there is nothing to check without fetching, and WebLLM's loader reports a bad base URL +far better than a HEAD request would. + +### How local registration works + +WebLLM composes every artifact URL as `new URL(relative, base)` and runs the base through `cleanModelUrl`, +which itself calls `new URL(...)` — so the base must be absolute and resolvable. A `blob:` URL cannot serve +as one, and there is no hook to hand the loader bytes directly. Pre-populating the cache under WebLLM's own +scopes and keys therefore **is** its native path: the loader does its ordinary thing and finds everything +already present. + +WebLLM is never told the model is local. Each model gets a synthetic base URL +(`https://local-model.invalid//resolve/main/`) and its artifacts are written into the exact cache +scopes and keys WebLLM's loader looks up: + +| Cache scope | Keys | +| --- | --- | +| `webllm/config` | `mlc-chat-config.json` | +| `webllm/model` | `tensor-cache.json`, tokenizer, every `params_shard_*.bin` | +| `webllm/wasm` | `-webgpu.wasm` | + +`reload()` therefore finds a full cache and issues zero requests. `.invalid` is reserved by RFC 6761 and can +never resolve — so this is the *mechanism* of the offline guarantee, not a label for it: there is no bug, no +eviction and no future refactor by which a locally-registered model reaches the network. It fails with a DNS +error instead. `test/sources.test.mjs` asserts that structurally, checking that every URL such a record +carries is on a `.invalid` host. + +`ModelStore.verify()` checks every key before a load, so storage eviction is reported as "re-register the +folder" rather than a mid-load fetch. It gates **only** the local route: a remote or prebuilt model that +loses its cache just re-downloads, which is slow, not fatal. + +`test/integration.test.mjs` pins this contract, including a guard that fails if a WebLLM upgrade renames a +cache scope or artifact. [test/e2e/run.mjs](test/e2e/run.mjs) proves it against a real model on a real GPU: +it temporarily wires a self-test page into the extension, drives ingest -> load -> streaming generation +through the production code paths, and restores the tree afterwards. Re-run it after bumping +`@mlc-ai/web-llm`. + +## Compiling a model in-house + +Done, for [`empero-ai/Qwen3.8-2B-Distill`](https://huggingface.co/empero-ai/Qwen3.8-2B-Distill). The full +record — every command, every toolchain breakage, and the numbers — is in +[MLC-COMPILE.md](MLC-COMPILE.md); the scripts are in `tools/`. Start with: + +```sh +tools/setup-mlc-toolchain.sh # venv, patches, emsdk, tvm web runtime +``` + +**Do not follow the upstream MLC quickstart.** No published pair of MLC nightlies works together, and five +separate breakages sit between `pip install` and a loadable `.wasm` — including two that produce a model +that compiles, ships and ingests cleanly and only fails at load. `setup-mlc-toolchain.sh` and +`patch-mlc-nightly.py` encode all of them and are idempotent; re-run after any `pip install`. + +**Why this model.** The text architecture is the one already verified end to end — two dimensions differ and +nothing else: + +| | Qwen3.5-0.8B (baseline) | Qwen3.8-2B-Distill (shipped) | +| --- | --- | --- | +| `model_type` | `qwen3_5` | `qwen3_5` | +| layers / heads / kv-heads / head_dim | 24 / 8 / 2 / 256 | 24 / 8 / 2 / 256 | +| `vocab_size`, `tie_word_embeddings` | 248320, true | 248320, true | +| `linear_*`, `full_attention_interval` | 16/128/16/128/4, every 4th | identical | +| `hidden_size` | 1024 | **2048** | +| `intermediate_size` | 3584 | **6144** | + +The wider `hidden_size` also helps: the GEMV reduction is split across 64 lanes, so 1024 leaves 2 iterations +per thread while 2048 leaves 4 and the 6144 `down_proj` leaves 12 — further along the measured efficiency +curve (x2 = 18.6, x8 = 27.0, x32 = 34.0 GB/s) before any schedule retune. + +**Quantization is not a tuning knob here.** Keep `q4f16_1`: unpacking eight nibbles and scaling them measured +~free (41 GB/s with the dequant arithmetic vs 41-46 without), so a wider format only adds bytes, and bytes are +exactly what decode pays for. `q3f16_1` is the only faster option (~1.28x) and it costs accuracy — less than +the 1.83x the schedule retune offers, and unlike it, not free. + +### Verifying a build + +```sh +node tools/audit-wasm.mjs # storage buffers per kernel +node tools/wasm-imports.mjs /*.wasm /*.wasm # runtime vs JS glue +MODEL_DIR= ENGINE_COUNT=2 npm run e2e +``` + +The import check is not optional: the model library links against whatever TVM web runtime built it but runs +against whatever tvmjs `@mlc-ai/web-llm` bundles, and a mismatch fails only at load, after everything else +has passed. Read the e2e's `decode probe` line too — `kernels/tok ÷ 24 layers` still ≈27 means dlight chose +more reduction threads over more work per thread and the schedule retune did not take. + +## Scheduling + +The engine is one GPU shared by every caller, so requests carry scheduling metadata and the engine — not the +caller — decides what runs when. Four mechanisms, no more ([src/background/pool.js](src/background/pool.js)): + +| | | +| --- | --- | +| **Priority bands** | `interactive` > `normal` (default) > `background`, FIFO within a band. | +| **Session supersession** | A new request with the same `session` cancels the previous one. This is the ghost-text primitive: each keystroke replaces the in-flight request instead of queueing behind it. | +| **Opt-in preemption** | An `interactive` request with no free engine interrupts a running job that set `preemptible: true`. The victim resolves with its partial output and is never requeued, so nothing can starve. | +| **One task, one engine** | Every request belongs to a `task` — a whole `batch` is one task, an unlabelled `chat` is its own. A task holds at most one engine, so two runnable tasks always run side by side whenever two engines exist. | + +Nothing else interrupts work in flight. A job that did not opt in always runs to completion. + +### Why one engine per task + +A batch used to spread across the whole pool. On the 0.8B that was worth 1.3-2.0x, because decode was +**sync-bound** — a stream spent most of its time waiting on Firefox's ~100 ms tick, so a second stream filled +idle GPU. Multi-step decoding removed most of that wait, and on the 2B the remaining cost is real GPU work. +Measured, same four prompts, greedy so both runs emit exactly 97 tokens: + +| pool | wall | aggregate | per item | peak overlap | +| --- | --- | --- | --- | --- | +| 1 | 8.2 s | 11.9 tok/s | ~2.0 s | 1 | +| 2 | 7.7 s | 12.6 tok/s | ~3.9 s | 2 | + +**1.06x.** The overlap is real — busy time summed to 15.2 s against 7.7 s of wall — but each stream runs at +half speed, so they cancel. Spreading one task over the pool buys ~nothing and costs the thing a second +engine is actually for: a page translation would sit on both engines while ghost-text waited behind it. + +So the rule is flat. An engine may idle while one task still has work queued; that ~6% is deliberately given +up to keep an engine free for whoever shows up next. `engine scaling:` in the e2e prints this ratio — +re-measure it per model, because on a small enough model the old fan-out logic would win again. + +### The pool grows, it is not sized + +`engineCount` is a **cap, not a size**. `load()` brings up exactly one engine; the pool adds another only +when a task that owns no engine is waiting. An engine no second task ever needed is ~1.6 GB on the 0.8B and +~2.4 GB on the 2B, bought for nothing. + +**There is no budget to check first.** Firefox implements neither `navigator.deviceMemory` nor +`performance.memory`, and `navigator.storage.estimate()` reports disk quota, not RAM — verified against the +shipped binary, and re-checked every run by the `memory signals:` line in the e2e. Nothing tells an extension +how much memory is left. So the pool does not predict, it probes: **a failed load is the memory check.** +Growth then stops for that model and is not retried, and `status().growthBlocked` says so. + +**Growth is not instant.** Building an engine is a full model load — measured **51 s** for the 2B — so both +tasks that triggered it will have finished first. The pool pays that once, in the background, and the second +engine is there for the *next* collision. If a workload is known to be concurrent from the start, the honest +fix is a warm-up request pair right after load, not a lower growth threshold. + +### Why each engine is a Web Worker + +Several MLCEngines cannot share a realm. Running the same e2e three ways isolates it: + +| setup | result | +| --- | --- | +| 1 engine, background page | passes | +| 2 engines, background page | both load, the first generates fine, the second's first generation fails: `Expected null or instance of VectorInt, got an instance of VectorInt` | +| 2 engines, one worker each | passes | + +The trigger is a second engine *generating* in the same realm — not the pool, and not the engine count by +itself. That message is embind reporting a type-registry mismatch, and the bundle does carry module-scoped +emscripten state (`var Module`, `var __wasmLib`) shared by every instance, which fits; but the fix rests on +the isolation above rather than on having traced the registry. + +Workers are viable because Firefox exposes WebGPU to dedicated workers and the 100 ms completion tick is +shared across them, so concurrency survives the move off the main thread (measured: 4 workers, 36.3 syncs/s). + +### How many engines are worth it + +Measured on the 0.8B, four independent prompts, back when fan-out still scaled: + +| pool | peak overlap | aggregate vs. serial | +| --- | --- | --- | +| 1 | 1 | 1.00x | +| 2 | 2 | 1.3x - 2.0x | +| 4 | 4 | **0.3x** — 3x slower than one engine | + +Four engines overlap in wall-clock terms but each drops from ~9.6 to ~0.7 tok/s: four copies of the weights +plus their KV caches leave a 16 GB machine with nothing free, and they starve each other. **More engines is +not more throughput** — and on the 2B, per the table above, a second one is not more throughput either. The +route to concurrent throughput is batched decode inside one engine (see Current Tasks), not more engines. + +## API + +### In-process — the library + +Migrating off `@mlc-ai/web-llm` costs one line; everything after it is unchanged. + +```js +-import { CreateMLCEngine } from "@mlc-ai/web-llm"; +-const engine = await CreateMLCEngine(modelId, { initProgressCallback }); ++import { CreateScheduledEngine } from "everything-webgpu"; ++const engine = await CreateScheduledEngine(modelId, { initProgressCallback }); + +await engine.chat.completions.create({ + messages, stream: true, + session: "ghost-text", // added — supersedes the previous request + priority: "interactive", // added — may preempt work that opted in +}); +``` + +`chat.completions.create()` returns WebLLM's own shapes, including its +`"stop" | "length" | "abort"` finish reasons. What it has no room for is `cancelled` and `preempted` +as distinct outcomes — both collapse to `"abort"` — which is why `complete()` below stays the direct +API rather than a legacy one. + +The fuller surface, when you want the store, the model source, or those outcomes: + +```js +import { ScheduledEngine, ModelStore } from "everything-webgpu"; +import { indexedDBStorage, ensurePersistent } from "everything-webgpu/adapters/idb"; + +// A page origin holds weights in *evictable* storage until this is granted. +await ensurePersistent(); + +const engine = new ScheduledEngine({ store: new ModelStore(await indexedDBStorage()) }); +await engine.load("Llama-3.2-1B-Instruct-q4f16_1-MLC"); + +const { text } = await engine.complete({ + messages: [{ role: "user", content: "hi" }], + session: "ghost-text", + priority: "interactive", +}); +``` + +| method | what it does | +| --- | --- | +| `load(id, {keepResident, signal})` | bring a model up; `signal` aborts the download, partial shards kept for a free resume | +| `use(id)` | switch between **resident** models — free, no reload | +| `unload(id?)` / `unloadAll()` | free VRAM, **keep the cached bytes** | +| `resident` | model ids with a live pool right now | +| `store.evict(id)` | free the disk, **keep the record** so it can be re-fetched | +| `remove(id)` | forget it entirely — frees bytes for **every** source, then drops the record | +| `chat.completions.create` | streamed chunks are WebLLM's own, verbatim: `tool_calls`, `logprobs`, stable `created` | +| `store.cacheState(rec)` | `"cached"` / `"partial"` / `"absent"` | +| `estimateSpeed(id?)` | projected tok/s, measured once anything has decoded | +| `features()` | what is actually switched on: KV reuse, decode steps, engines | +| `complete(req, onChunk?)` | one completion; `onChunk` streams deltas | +| `batch(req, onItem?)` | many independent prompts as **one task** — see below | +| `cancel(idOrSession)` | by job id or session key; returns how many it stopped | +| `configure({ decodeSteps })` | retune a live engine, no reload | +| `registerModel(spec)` | a base URL, or local `files` | +| `listModels()` / `listAvailableModels()` | registered only (cheap) / all three routes | +| `subscribe(fn)` | lifecycle changes; returns an unsubscribe | +| `state` / `hasWebGPU` | current snapshot, WebGPU presence | +| `chat.completions.create(req)` | the WebLLM/OpenAI facade over `complete()` | +| `probe()` | WebGPU, adapter, `shader-f16`, limits, storage quota — cached | +| `canRun(modelId)` | `{ ok, blockers, warnings }`, before anything is downloaded | +| `recommendModels({maxVramMB, prefer})` | rank the 163 prebuilt models for *this* device | + +Failures are `EngineError { code, message, detail }` — `NO_WEBGPU`, `NO_MODEL`, `UNKNOWN_MODEL`, +`CACHE_INCOMPLETE`, `INVALID_MODEL_FOLDER`, `BAD_REQUEST`, `GENERATION_FAILED`. `detail` carries the +structured context (the evicted keys, the missing field, why a folder was rejected), so no caller +parses a message. Over the wire the code rides beside `error`, which stays a plain string. + +Generation ops take `messages`, `temperature`, `max_tokens`, `response_format` and `extra_body` — the +OpenAI shape WebLLM already speaks — plus the scheduling fields `task`, `session`, `priority` and +`preemptible`, which are what this adds over calling WebLLM directly. + +### Over a wire — the WebExtension adapter + +Only for the case where the engine and the caller are in different processes. Extension id: +`everything-webgpu@local`; the manager page prints a copy-pasteable version. + +Two transports, one vocabulary ([src/adapters/protocol.js](src/adapters/protocol.js)): + +- `browser.runtime.sendMessage(id, req)` — request/response. Ops: `status`, `listModels`, `load`, + `unload`, `chat`, `batch`, `cancel`, `configure`. +- `browser.runtime.connect(id, { name: "everything-webgpu/v1" })` — streaming. Ops: `subscribe`, + `chat.stream`, `batch.stream`, `abort`; the port also pushes `engineState` on every lifecycle change. + +Every message carries `protocol`, so a stray message from another sender fails fast instead of being +half-interpreted. Generation ops accept `modelId`, `temperature`, `max_tokens`, `response_format`, +`extra_body`, plus the scheduling fields `task`, `session`, `priority` and `preemptible`. + +**Send raw requests. Do not ask for a `translate` op.** The engine schedules a shared GPU; it does not +author prompts. Prompts belong to whoever owns the feature, because they are model-specific — switching this +build from `Qwen3.5-0.8B` to `Qwen3.8-2B-Distill` changed the conversation template and made every reply open +with a `` block. A prompt that lives in the caller survives that; a `translate` op baked into the +engine would have to be rewritten and re-shipped to every caller. Wrap the transport in a client-side helper +if you want `translate()` ergonomics — just keep it on your side of `sendMessage`. + +### The three shapes of work + +What differs between these is *not* the op or the transport. It is who owns an engine, and what may +interrupt what. The examples below use the wire form; in-process the same fields go to `complete()` and +`batch()`. + +| | op | priority | key fields | why | +| --- | --- | --- | --- | --- | +| **Completion** (ghost text) | `chat.stream` | `interactive` | `session` | Each keystroke supersedes the last request; may preempt opted-in work. | +| **Translation** (a page) | `batch` | `normal` | one shared `task` | One request instead of N, so the engine schedules it as a unit and it never hogs the pool. | +| **Reformat** (markdown) | `chat` | `background` | `preemptible: true` | Nobody is watching; let interactive work cut in. | + +#### Completion — latency is the whole product + +```js +const port = browser.runtime.connect("everything-webgpu@local", { name: "everything-webgpu/v1" }); +port.onMessage.addListener((m) => { + if (m.op === "chunk") render(m.delta); + if (m.op === "done") finish(m.text); +}); + +// On every keystroke. The previous request is cancelled, not queued behind. +port.postMessage({ + protocol: "everything-webgpu/v1", + op: "chat.stream", + id: crypto.randomUUID(), + session: "ghost-text", // supersession key — the important field + priority: "interactive", // may preempt jobs that opted in + max_tokens: 24, // ghost text is short; do not pay for more + messages: [{ role: "user", content: prefix }], +}); +``` + +`session` is what makes this work, not `cancel`. Reusing one session key means the engine drops the stale +request itself; a caller that mints a fresh id per keystroke and calls `cancel` races its own typing. + +#### Translation — throughput, one task + +```js +// One batch, not a loop of `chat` calls. +const res = await browser.runtime.sendMessage("everything-webgpu@local", { + protocol: "everything-webgpu/v1", + op: "batch", + task: "translate-page", // optional; a batch is one task either way + requests: sentences.map((s) => ({ + messages: [{ role: "user", content: `Translate to French, output only the translation:\n${s}` }], + })), +}); +res.results.forEach((r) => apply(r.index, r.text)); +``` + +Every item of one batch shares a task, and a task holds one engine, so a 200-sentence page occupies exactly +one engine and can never freeze ghost-text behind it. Results carry `engineIndex`, `startedAt` and +`finishedAt`, so a caller can check what actually ran where. + +`batch` is still the right call rather than a loop of `chat`: it is one round trip, the engine keeps the +items in one queue it can reason about, and if this ever runs on a model small enough for fan-out to pay +again — or once batched decode lands — the same call gets faster with no change on your side. + +Use `batch.stream` over a port instead if you want items as they land rather than one array at the end. + +#### Reformat — cheap to interrupt + +```js +await browser.runtime.sendMessage("everything-webgpu@local", { + protocol: "everything-webgpu/v1", + op: "chat", + priority: "background", + preemptible: true, // the direction matters — see below + max_tokens: 2048, + messages: [{ role: "user", content: `Reformat as clean Markdown, no commentary:\n\n${doc}` }], +}); +``` + +**Set `preemptible` on the work that can afford to lose, not on the work you care about.** Only an +`interactive` request preempts, and only a job that opted in can be preempted. A preempted job resolves with +`preempted: true` and whatever text it had, so it is never requeued and can never starve — but that also +means you must be able to use, or discard, a partial result. + +### Getting these wrong + +| symptom | cause | +| --- | --- | +| Ghost text lags behind typing | Fresh `id` per keystroke with no `session`, so every stale request still runs. | +| Page translation is slower than expected | Expected: one task is one engine, and a second engine measured 1.06x anyway. Throughput here comes from batched decode, not from more engines. | +| Reformatting blocks completions | `preemptible` left off the background job, so `interactive` has nothing to take. | +| Pool stays at one engine | Expected: it grows only when a *second task* waits. Check `status().growthBlocked` if two are waiting and it still has not. | + +By default every installed extension may call the API. The manager page has an allowlist field; fill it in +with extension ids to restrict access. + +## Layout + +| Path | Role | +| --- | --- | +| [manifest.json](manifest.json) | MV2, persistent background page, `wasm-unsafe-eval` CSP | +| [src/engine/index.js](src/engine/index.js) | Public entry point of the library | +| [src/engine/create.js](src/engine/create.js) | `CreateScheduledEngine` — the one-line swap for `CreateMLCEngine` | +| [src/engine/chat.js](src/engine/chat.js) | `chat.completions.create()`, the WebLLM/OpenAI facade | +| [src/engine/environment.js](src/engine/environment.js) | `environment()` — the read-only device/runtime report; writes are `configure()` | +| [src/engine/errors.js](src/engine/errors.js) | `EngineError` and the eight codes | +| [src/engine/device.js](src/engine/device.js) | Hardware probe, `canRun`, model ranking | +| [src/engine/engine.js](src/engine/engine.js) | `ScheduledEngine` — the engine with no transport attached | +| [src/engine/pool.js](src/engine/pool.js) | Engine pool + priority scheduler | +| [src/engine/engine-worker.js](src/engine/engine-worker.js) | One pool slot's engine, in its own realm | +| [src/engine/multistep.js](src/engine/multistep.js) | Multi-step decoding: K forward steps per GPU sync | +| [src/engine/sources.js](src/engine/sources.js) | What `load()` was handed — pure dispatch across id / URL / spec / folder | +| [src/engine/ingest.js](src/engine/ingest.js) | Folder validation and cache injection | +| [src/engine/recipes.js](src/engine/recipes.js) | `ask()` / `conversation()` / `ghostText()` — the three shapes as one call each, scheduling only | +| [src/engine/prefetch.js](src/engine/prefetch.js) | `prefetch()` — fill the cache with no engine and no GPU; WebLLM's `hasModelInCache` is the oracle | +| [src/engine/model-store.js](src/engine/model-store.js) | Cache layout, registry, settings, `StorageAdapter`, the three model sources | +| [src/engine/constants.js](src/engine/constants.js) | `PRIORITY`, `ENGINE_STATE` — engine vocabulary, transport-free | +| [src/adapters/protocol.js](src/adapters/protocol.js) | Wire protocol: `PROTOCOL`, `OP`, `PORT_OP` | +| [src/adapters/webext.js](src/adapters/webext.js) | `browser.storage.local` + the message/port router | +| [src/adapters/idb.js](src/adapters/idb.js) | IndexedDB `StorageAdapter` + `ensurePersistent()`, for pages | +| [src/adapters/memory.js](src/adapters/memory.js) | In-memory `StorageAdapter`, for tests | +| [src/background/background.js](src/background/background.js) | The extension host: build an engine, attach the transport | +| [src/popup/](src/popup/) | Minimal test chat (moves to `demo` in Phase 4) | +| [src/manager/](src/manager/) | Drop target, registry, settings, setup help (moves to `demo` in Phase 4) | +| [test/integration.test.mjs](test/integration.test.mjs) | The cache-injection contract | +| [test/scheduler.test.mjs](test/scheduler.test.mjs) | Priority, supersession, preemption and pool growth, GPU-free | +| [test/sources.test.mjs](test/sources.test.mjs) | How `load()` resolves prebuilt / remote / injected, in what order it refuses, and that a local model has no reachable URL | +| [test/errors.test.mjs](test/errors.test.mjs) | That failures carry the right code, and that nothing throws an untyped Error | +| [test/chat.test.mjs](test/chat.test.mjs) | That the WebLLM facade really is drop-in, shape by shape | +| [test/device.test.mjs](test/device.test.mjs) | The compatibility rules, and that blockers and warnings stay distinct | +| [test/manage.test.mjs](test/manage.test.mjs) | The four model states, and that unload / evict / remove stay distinct | +| [test/e2e/](test/e2e/) | Real-hardware end-to-end run (`npm run e2e`) | +| [test/e2e/bench.mjs](test/e2e/bench.mjs) | Standalone WebGPU sync-latency benchmark (`npm run bench`) | +| [MLC-COMPILE.md](MLC-COMPILE.md) | How the model was compiled, and every toolchain breakage on the way | +| [tools/](tools/) | Model-compilation toolchain: setup, nightly patches, weight strip, wasm audits | +| [WEBLLM-SURFACE.md](WEBLLM-SURFACE.md) | What WebLLM already does, what we add, and where the line is. **Read before adding a capability**; its "Upgrading" section is the dependency-bump runbook. | +| [ROADMAP.md](ROADMAP.md) | The only list of open work | +| [ARCHIVE.md](ARCHIVE.md) | What was done and why — the extraction, the model sources, the de-duplication | + +`src/engine/` references no WebExtension API — asserted by a test, because that claim is only broken in the +host nobody ran. The three places the host used to leak in are injected: a `StorageAdapter` for the registry, +a worker URL, and the WebLLM import. `src/adapters/` holds one implementation of each per host. See +[ARCHIVE.md](ARCHIVE.md). + +In *this* host the engine lives in the MV2 persistent background page — a real document on the extension +origin, so it has both `navigator.gpu` and the same Cache Storage the manager page writes to. The model stays +resident in VRAM across popup opens and across calls from other extensions. + +## Known limits + +- **AMO signing**: `vendor/web-llm.js` is ~6 MB, over `web-ext lint`'s 5 MB parse limit. Fine for temporary + install and self-distribution; it would need splitting before an AMO listing. +- **MV2**: MV3 event pages get evicted, which would unload a multi-GB model between calls. Migrating needs a + keep-alive or a dedicated engine tab. This constrains `demo` only — the library is host-agnostic. +- **Storage eviction off an extension origin**: an ordinary page has no `unlimitedStorage`, so a multi-GB + model is evictable until `ensurePersistent()` is granted. For a prebuilt or remote model that means a slow + reload; for a locally-registered one it is fatal and it must be re-registered. Unmeasured: the exact quota + and grant behaviour per browser. +- **Never run outside a Firefox extension**: the library is written against capabilities rather than + browsers, and should be *faster* on Chrome (KV reuse is not disabled there — see "Multi-step decoding"). + Both claims are predictions. See ROADMAP.md, Gates A and B. +- **Thinking burns the budget**: this model opens every reply with a `` block — it is a reasoning + distill and its card says so. At ~17 tok/s that is seconds of nothing before the answer starts. For + translation and completion, suppress it in your prompt, or rebuild the config with the `qwen3_5_nothink` + conversation template. The engine deliberately does not rewrite prompts for you. +- **No cross-turn KV reuse**: every turn re-prefills the whole history at 5.27 ms/token, so a long + conversation pays ~22 s before its first token at the 4096 limit. See Current Tasks. +- **Decode speed**: *was* bounded by Firefox's 100 ms WebGPU poll timer (Mozilla bug 1870699). Multi-step + decoding and compute-pass batching bought most of that back; what remains is memory bandwidth — decode + achieves ~16 GB/s of the M4's ~120. See "The 10 tok/s ceiling" and "Why not llama.cpp/Ollama-class". +- **A second engine is not more speed**: measured 1.06x on this model. It buys isolation between tasks. + Concurrent throughput needs batched decode inside one engine. See "Scheduling". diff --git a/API.md b/API.md new file mode 100644 index 0000000000000000000000000000000000000000..dd80908479e53ff3f54fa49065554594635249bb --- /dev/null +++ b/API.md @@ -0,0 +1,367 @@ +# API — every way to call it + +The complete call surface of `everything-webgpu`, one page. [README.md](README.md) is the pitch and +the migration story; this is the catalogue. Asserted against the code by +[test/api-doc.test.mjs](test/api-doc.test.mjs) — every engine method named here exists, every +package export appears, and the error table equals `ERROR`, so this page cannot drift from the +source without failing `npm test`. + +The page is in three tiers. **[Start here](#start-here)** — `load`, `ask`, `conversation`, +`environment` — is the dead-simple path, and most apps need nothing else. **[Native +passthrough](#native-passthrough)** is `chat.completions.create()`, the WebLLM/OpenAI compatibility +layer, which **never changes** (see [Stability](README.md#stability)). **[When you need +more](#when-you-need-more)** is the rest of the surface: the general calls and their scheduling +fields, embeddings, residency, device inspection, configuration, lifecycle. Everything outside the +passthrough is pre-1.0 and may move; the ergonomic verbs are being consolidated in +[ROADMAP.md](ROADMAP.md). + +--- + +## The four lines + +```js +import { CreateScheduledEngine } from "everything-webgpu"; // 1. import +const engine = await CreateScheduledEngine("Llama-3.2-1B-Instruct-q4f16_1-MLC"); // 2. load a model +const reply = await engine.ask("Name three primary colours."); // 3. ask +console.log(reply); // 4. the answer +``` + +`reply` is a plain string. Line 2 downloads ~0.8 GB the first time and prints throttled progress to +the console unless you pass `initProgressCallback` (render it yourself) or `initProgressCallback: +null` (silence). After that first run it is a cache read and needs no network. + +On Vite, add one plugin — see [Bundlers](#bundlers). + +--- + +## Start here + +Four calls. Get an engine with `CreateScheduledEngine` (below, under [Getting an +engine](#getting-an-engine)), then `load` a model, `ask` it things or hold a `conversation`, and +`environment` tells you whether the machine is up to it. + +### Load a model — `engine.load(src, opts?)` + +**One call, four source shapes.** `engine.load(src, opts?)` works out what you handed it, registers +whatever needs registering, and brings the model up. + +| `src` | route | +| --- | --- | +| `"Llama-3.2-1B-Instruct-q4f16_1-MLC"` | a prebuilt id, or anything you registered earlier. A typo is answered with near matches. | +| `"https://huggingface.co/mlc-ai/Foo-MLC"` | an HF repo. `/resolve/main/` is **not** derived — WebLLM appends it. | +| `"https://cdn.example/models/foo/"` + `{ modelLib }` | any base URL you host. `modelLib` is **required** and never guessed (0 of 163 prebuilt models have a derivable lib name or same-origin lib). | +| `{ model, modelLib }` | the explicit remote spec. | +| `input.files` \| `dropEvent.dataTransfer` \| `{ files }` | a folder off disk. **No network at any point.** | + +**`opts`** — `keepResident`, `signal`, `modelType`, `contextWindow`, `vramRequiredMB`, `id`, +`onProgress`, `defer`. + +- `id` overrides the id derived from the URL's last segment. +- `defer: true` registers the source **without building a pool** — the drop-now-load-later flow. It + returns the registry record. `defer` on a bare prebuilt id is an error, not a silent load. +- `keepResident: true` holds this model in VRAM alongside whatever is already up. The default + unloads everything else first — the safe choice on a 16 GB machine. + +```js +await engine.load("Llama-3.2-1B-Instruct-q4f16_1-MLC"); +await engine.load("https://cdn.example/models/my-model/", { modelLib: "https://cdn.example/models/my-model/lib.wasm" }); +await engine.load(dropEvent.dataTransfer); +await engine.load(input.files, { defer: true }); // register now, build the pool on first use +``` + +`load()` composes the lower-level `registerModel`, `ingestModelFolder` and the download primitives; +`prefetch()` warms the cache with no GPU. All of that is under [More on +loading](#more-on-loading). + +### Ask one question — `engine.ask(input, opts?)` + +One question, its own task, **no session** — two `ask()`s never supersede each other. Returns the +reply string. `opts.onDelta` to stream. Goes through the same scheduler as everything else: priority +bands, one-task-one-engine, opt-in preemption. + +### Hold a conversation — `engine.conversation(opts)` + +A multi-turn chat that keeps its own history. `engine.conversation({ system?, keep?, ...defaults })` +— one stable task for every turn, turns serialised, history bounded at `keep: 12` exchanges +(`Infinity` opts out). + +```js +const chat = engine.conversation({ system: "You are terse." }); +await chat.say("capital of France?"); // → { text, finishReason } +await chat.say("and its population?", onDelta); // remembers +chat.messages; chat.length; chat.reset(); chat.restore(messages); +``` + +### Inspect the machine — `engine.environment(opts?)` + +| call | answers | +| --- | --- | +| `await engine.environment(opts?)` | **the preflight.** A report; every line has `severity` · `affects` · `cause` · `fix` · `operable`. `{ scope: "local" }` never touches the model layer (poll freely); `{ scope: "device" }` is hardware only. | +| `await engine.environment.measure()` | one calibration generation → measured tok/s for the current model | + +**`environment()` only reports.** Writes go through `configure()`; passing a setting to +`environment()` is an error that names `configure()`. Per-model "will it run" is +[`canRun(modelId)`](#inspecting-the-machine); the fuller device surface is there too. + +## Native passthrough + +The one call that **never changes**. `@mlc-ai/web-llm` is OpenAI-shaped, and so is this: the +migration off it is a one-line import swap, and `chat.completions.create()` then takes and returns +exactly the same shapes — same streamed chunk objects, same finish reasons, same non-streaming +envelope. See [Stability](README.md#stability). + +| call | shape | +| --- | --- | +| `engine.chat.completions.create(params)` | the **WebLLM/OpenAI** shape, unchanged. Streams the same chunks, same finish reasons. `session`/`priority`/`task`/`preemptible` are additive. | + +**`params`** = the OpenAI generation fields WebLLM already speaks (`messages`, `temperature`, +`max_tokens`, `response_format`, `extra_body`, …) **plus** the scheduling fields that are the only +thing this adds over calling WebLLM directly: + +| field | meaning | +| --- | --- | +| `modelId` | load/route to this model instead of the current one | +| `id` | job id; also what `cancel(id)` takes | +| `task` | the unit that owns an engine; a whole batch shares one | +| `session` | a later job with this key supersedes the earlier one | +| `priority` | `"interactive"` \| `"normal"` \| `"background"` | +| `preemptible` | may be interrupted by an `interactive` job (set it on the work that can afford to lose) | + +`complete()`, `completeRaw()` and `batch()` take the same scheduling fields and expose `cancelled` / +`preempted` as first-class outcomes the OpenAI shape has no room for — see [The general +calls](#the-general-calls). + +## When you need more + +### Getting an engine + +| call | when | +| --- | --- | +| `await CreateScheduledEngine(modelId?, opts?)` | the common case. Loads `modelId` before returning, like WebLLM's `CreateMLCEngine`. Omit it for an engine that loads later. | +| `new ScheduledEngine({ store, workerUrl?, loadWebLLM?, prebuilt? })` | when you must pass a store explicitly — a worker, a test, an extension. Does **not** load anything. | + +**`opts` for `CreateScheduledEngine`** — `store`, `initProgressCallback`, `workerUrl`, `loadWebLLM`, +`prebuilt`. Anything not `store`/`initProgressCallback` is forwarded to the constructor. + +| constructor field | default | meaning | +| --- | --- | --- | +| `store` | IndexedDB (`CreateScheduledEngine` only; the constructor requires it) | a `ModelStore`, or a bare `StorageAdapter` it wraps | +| `prebuilt` | `true` | expose WebLLM's 163 HuggingFace models. `false` = offline-only: `load()` resolves registered models and nothing else, and an unknown id fails before the WebLLM bundle is fetched | +| `workerUrl` | `new URL("./engine-worker.js", import.meta.url)` | the decode worker's module URL | +| `loadWebLLM` | `() => import("../../vendor/web-llm.js")` | override the bundle source (tests) | + +**Stores** — `import { indexedDBStorage } from "everything-webgpu/adapters/idb"` (pages, plus +`ensurePersistent()`), `everything-webgpu/adapters/memory` (`memoryStorage()`, tests), +`everything-webgpu/adapters/webext` (`webExtensionStorage()` + `attachWebExtensionTransport()`). + +```js +import { ScheduledEngine, ModelStore } from "everything-webgpu"; +import { memoryStorage } from "everything-webgpu/adapters/memory"; +const engine = new ScheduledEngine({ store: new ModelStore(memoryStorage()) }); +``` + +### More on loading + +**Warming the cache first** — `await engine.prefetch(modelId, { onProgress, signal })`. Downloads +the weights **without building an engine and without WebGPU**, so an app can warm the cache before +it knows whether the machine can run the model. Interrupted downloads resume; a second call is free. +WebLLM cannot express this — `reload()` needs a GPU before it fetches a shard. + +**Low-level, still exported** — `load()` composes these rather than replacing them: + +| call | does | +| --- | --- | +| `engine.registerModel(spec)` | add a `{ modelId, model, modelLib }` or `{ modelId, files }` record, no pool | +| `ingestModelFolder(entries, { store })` | folder → populated Cache Storage, returns the record | +| `filesFromInput(input.files)` / `filesFromDataTransfer(dt)` | either browser shape → flat `{ path, file }[]` (the latter is async) | +| `prefetchModel({ modelId, record, ... })` / `resolveModelUrl(...)` | the engine-free download primitives | + +### The general calls + +Every call here goes through the **same scheduler** as `ask` / `conversation` / `chat.completions`: +priority bands, `session` supersession, one-task-one-engine, opt-in preemption. + +| call | shape | +| --- | --- | +| `await engine.complete(payload, onChunk?)` | `{ text, usage, finishReason, cancelled?, preempted? }`. `onChunk(delta)` streams plain text. | +| `await engine.completeRaw(payload, onRawChunk?)` | same, but `onRawChunk` gets WebLLM's chunk object verbatim. | +| `await engine.batch({ requests, task?, ...sched }, onItem?)` | `requests` fanned across the pool as **one task**. Returns `BatchItem[]` — each with `index`, `engineIndex`, `startedAt`, `finishedAt`, and `text`/`usage` or `error`. | + +**`payload`** is the OpenAI generation fields plus the scheduling fields — the same table as [Native +passthrough](#native-passthrough) (`modelId`, `id`, `task`, `session`, `priority`, `preemptible`). + +Result flags: `cancelled: true` (superseded or `cancel()`ed), `preempted: true` (`text` is partial). + +### Ghost text — `engine.ghostText(opts)` + +`engine.ghostText({ prompt, debounceMs?, maxTokens?, session?, ...defaults })` — debounce + one +session key + `interactive` priority + **resolves `null` when stale**. `prompt` is **required**, no +default: prompts are model-specific and belong to whoever owns the feature. + +```js +const ghost = engine.ghostText({ prompt: (before) => `Continue:\n${before}` }); +const hint = await ghost.suggest(editor.textBefore()); // string | null +ghost.cancel(); // on blur / accept +``` + +### Embeddings + +Needs an **embedding model** (`snowflake-arctic-embed-*`, from 239 MB), usually held resident +alongside a chat model. + +| call | returns | +| --- | --- | +| `await engine.embed(input, opts?)` | `number[][]` — one vector per input, in order | +| `await engine.embedRaw(input, opts?)` | WebLLM's OpenAI envelope (`data[].embedding`) | + +`opts`: `modelId`, `task`, `session`, `priority`, `preemptible`, `id`. **A running embedding cannot +be interrupted** — one forward pass has no decode loop to break out of; queued embeddings supersede +normally. + +### Residency and cache + +A resident model is a full copy of its weights in VRAM, and nothing reports free VRAM to a page — +so residency is explicit. + +| call | frees | keeps | +| --- | --- | --- | +| `await engine.unload()` | current model's VRAM | cache + registry | +| `await engine.unload(id)` | that model's VRAM | cache + registry | +| `await engine.unload(id, "cache")` | VRAM + cached bytes | registry entry | +| `await engine.unloadAll()` | every resident model's VRAM | cache + registry | +| `await engine.remove(id)` | bytes + registry entry | nothing — for an injected model, means re-supplying the folder | +| `engine.evict(id)` | low-level primitive `unload(id, "cache")` is built on | registry entry | + +**Routing without loading** — `engine.use(id)` points unaddressed requests at an already-resident +model (free and instant; `load()` is what costs). `engine.resident` lists model ids with a live +pool. `await engine.cacheState(id)` says what is on disk (`"complete"` / `"partial"` / absent). + +### Inspecting the machine + +The device surface behind the [Start here](#start-here) preflight. + +| call | answers | +| --- | --- | +| `await engine.canRun(modelId)` | per-**model**: `{ ok, blockers, warnings }`, before anything downloads | +| `await engine.recommendModels(opts?)` | which models this device should be asked to run, best first. `opts`: `maxVramMB`, `needsVision`, `needsToolCalling`, `prefer` | +| `await engine.estimateSpeed(modelId?)` | projected decode tok/s (uses the measured rate once one generation has happened) | +| `await engine.probe()` | raw device probe: WebGPU, adapter, `shader-f16`, the five limits, storage quota. Cached. | +| `await engine.features()` | what is switched **on** now, vs what the device could support. `multiStepOff` is non-null when decode fell back to one GPU sync per token — the silent halving `environment()` reports as `degraded` | +| `engine.hasWebGPU` | `Boolean(navigator.gpu)` | +| `await engine.listAvailableModels()` | registered + prebuilt, normalised. Costs one bundle fetch. | +| `engine.listModels()` | registered only — cheap, no bundle load | + +### Configuration + +`await engine.configure(patch)` — applies a runtime knob and persists it as the default. + +| knob | effect | +| --- | --- | +| `decodeSteps` | forward steps per GPU sync. Hot, no reload. `1`–`32` (`DEFAULT_DECODE_STEPS` = 15). | +| `engineCount` | pool size. Persisted; live pools keep the size they came up with. | +| `temperature`, `maxTokens`, `systemPrompt` | generation defaults (`DEFAULT_SETTINGS`) | + +Not operable from JS, report-only via `environment()`: KV reuse (derived from the 9-storage-buffer +cap), compute-pass batching (build-time `NO_PASS_MERGE`), `shader-f16`, GPU, `about:config` flags. + +### Lifecycle and cancellation + +| call | | +| --- | --- | +| `const stop = engine.subscribe(listener)` | `listener(state)` fires immediately, then on every change. Returns unsubscribe. | +| `engine.state` | snapshot: `status`, `modelId`, `progress`, `error`, `pool {size,busy,queued,maxSize,growthBlocked}`, `resident`, `decode` | +| `engine.store` | the `ModelStore`, so a host can drive the registry without a second handle | +| `engine.cancel(idOrSession)` | cancel by job id or by session key | +| `engine.load(id, { signal })` | an `AbortController` signal tears down an in-flight download | + +`state.status` is one of `ENGINE_STATE`: `"idle"` · `"loading"` · `"ready"` · `"error"`. + +## Errors + +Every failure is an `EngineError` with a `.code`, a human-readable `.message` (the thing you print), +and structured `.detail`. `import { isEngineError, ERROR } from "everything-webgpu"`. + +| code | what to do | +| --- | --- | +| `NO_WEBGPU` | tell the user to check flags/hardware; retrying is futile | +| `NO_MODEL` | nothing registered — send them to your setup flow | +| `UNKNOWN_MODEL` | that id is not resolvable; `listAvailableModels()` says what is | +| `CACHE_INCOMPLETE` | a locally-registered model was evicted; re-register the folder | +| `INVALID_MODEL_FOLDER` | not a compiled MLC model; `detail` says what is missing | +| `BAD_REQUEST` | the caller's arguments are wrong — a bug in the caller | +| `ABORTED` | the caller cancelled it. Not a failure; do not report it as one | +| `GENERATION_FAILED` | the model failed mid-generation | +| `PACKAGE_INCOMPLETE` | your **build** is wrong, not your code — missing `vendor/` bundle, or a decode worker the bundler did not emit. `message` names the fix; `detail.cause` says which | + +```js +try { await engine.load(id); } +catch (err) { + if (isEngineError(err, ERROR.CACHE_INCOMPLETE)) return reRegisterFolder(); + throw err; +} +``` + +## Bundlers + +The engine spawns its decode worker with `new Worker(new URL("./engine-worker.js", +import.meta.url), { type: "module" })`. On **Vite**, its dependency pre-bundler rewrites that URL to +a path that 404s — in `vite dev`, on a real (non-linked) install only. Add the plugin: + +```js +import { everythingWebGPU } from "everything-webgpu/vite"; +export default defineConfig({ plugins: [everythingWebGPU()] }); +``` + +Equivalent by hand: `optimizeDeps: { exclude: ["everything-webgpu"] }`. Skip both and `load()` +throws `PACKAGE_INCOMPLETE` naming the fix rather than hanging. `vite build` is unaffected either +way. Other bundlers that honour `new URL(..., import.meta.url)` for workers (Webpack 5, Rollup, +Parcel 2) need nothing. + +## Every export + +`import { … } from "everything-webgpu"` — 43 names. + +**Engine & entry** — `ScheduledEngine`, `CreateScheduledEngine`, `EnginePool` + +**Model sources** — `ModelStore`, `ingestModelFolder`, `filesFromInput`, `filesFromDataTransfer`, +`prefetchModel`, `resolveModelUrl`, `isInjected`, `baseUrlFor`, `groupKeysByScope`, `toAppConfig` + +**Recipes** (also methods on the engine) — `ask`, `conversation`, `ghostText` + +**Device** — `probeDevice`, `canRun`, `projectSpeed`, `rankModels`, `REFERENCE_DECODE_BYTES_PER_SECOND` + +**Multi-step decoding** — `installMultiStepDecoding`, `burstSize`, `clampSteps`, +`DEFAULT_DECODE_STEPS`, `MAX_DECODE_STEPS` + +**Errors** — `EngineError`, `ERROR`, `isEngineError`, `asEngineError` + +**Formatting** — `formatBytes` + +**Enums / constants** — `PRIORITY`, `PRIORITY_ORDER`, `ENGINE_STATE`, `UNLOAD_LEVEL`, `SEVERITY`, +`MODEL_TYPE`, `SOURCE`, `DEFAULT_SETTINGS`, `WORKER_CONFIGURE`, `CACHE_CONFIG`, `CACHE_MODEL`, +`CACHE_WASM` + +### Enum values + +| enum | values | +| --- | --- | +| `PRIORITY` | `interactive` · `normal` · `background` | +| `ENGINE_STATE` | `idle` · `loading` · `ready` · `error` | +| `UNLOAD_LEVEL` | `vram` · `cache` | +| `SEVERITY` | `blocked` · `degraded` · `tune` · `info` · `ok` | +| `SOURCE` | `prebuilt` · `remote` · `injected` | +| `MODEL_TYPE` | `llm` = 0 · `embedding` = 1 · `vlm` = 2 | +| `CACHE_*` | `webllm/config` · `webllm/model` · `webllm/wasm` | +| `DEFAULT_SETTINGS` | `engineCount: 2`, `decodeSteps: 15`, `temperature: 0.6`, `maxTokens: 1024`, `systemPrompt: ""` | + +## Subpath exports + +| specifier | | +| --- | --- | +| `everything-webgpu` | everything above | +| `everything-webgpu/vite` | `everythingWebGPU()` Vite plugin | +| `everything-webgpu/worker` | the decode worker entry (for a custom `workerUrl`) | +| `everything-webgpu/adapters/idb` | `indexedDBStorage()` | +| `everything-webgpu/adapters/memory` | `memoryStorage()` | +| `everything-webgpu/adapters/webext` | `webExtensionStorage()`, `attachWebExtensionTransport()` | +| `everything-webgpu/adapters/protocol` | the wire-protocol constants | diff --git a/ARCHIVE.md b/ARCHIVE.md new file mode 100644 index 0000000000000000000000000000000000000000..3a7832698e540d3ed8471d445cb8fb2ae0cd83a3 --- /dev/null +++ b/ARCHIVE.md @@ -0,0 +1,493 @@ +# Archive — done, and why + +Completed work, kept for its **reasoning** rather than its status. A decision recorded without its +cause gets re-litigated by the next session, or quietly reversed. + +Open work is in [ROADMAP.md](ROADMAP.md). What is true and measured is in [AI.md](AI.md). + +--- + +## The extraction — `main` became a library + +`main` is a library for developers embedding a local model in their own app; `demo` keeps the Firefox +extension and is the library's first consumer. + +**Phase 0 — froze the extension.** `demo` branch, `demo-baseline` tag. + +**Phase 1 — decoupled the three platform dependencies.** `background.js` went from 390 lines to 22. +The engine core (`src/engine/`) references no WebExtension API, asserted by a test — that claim is +only ever broken in the host nobody ran. + +| | before | after | +| --- | --- | --- | +| Registry | `browser.storage.local` ×4 | `ModelStore` over an injected `StorageAdapter` — deliberately `get(key)`/`set(obj)`, the exact shape of `browser.storage.local`, so the WebExtension adapter is a passthrough | +| Worker | `browser.runtime.getURL(…)` | `new URL("./engine-worker.js", import.meta.url)` — understood by Vite/webpack/esbuild *and* correct on `moz-extension://`, so it replaces `getURL` rather than sitting beside it | +| Transport | 6 `browser.runtime` listeners | `attachWebExtensionTransport(engine)`; wire format unchanged byte-for-byte | + +The wire protocol was **demoted**, not removed: in a page there is no `sendMessage`, so `OP`/`PORT_OP` +became one adapter's vocabulary rather than the interface. + +**Phase 2 — the developer-facing surface.** + +- `CreateScheduledEngine()` + `chat.completions.create()`. Migration off `@mlc-ai/web-llm` costs one + line; everything after it is unchanged. `complete()`/`batch()` keep the names so `chat` stays free + for the facade. +- `EngineError { code, message, detail }` — 8 codes, each existing because a caller does something + *different* about it. A test walks `src/engine/` for bare `throw new Error`, because one untyped + throw forces every caller back to string matching. +- `.d.ts` generated from the JSDoc that was already there; only the request/result typedefs were + added. Source stays plain ESM. +- `@mlc-ai/web-llm` pinned to exactly `0.2.84` — `build.mjs` patches the bundle by string anchor and + throws on a miss, so a caret range would break every consumer's build. + +## Zero-download stopped being a constraint + +It existed because the *extension* could not reasonably download. A developer embedding a model +usually can, and often must. So a model now arrives by one of three routes, resolved by `load()`: + +| route | how | validation | +| --- | --- | --- | +| `prebuilt` | one of WebLLM's 163 HuggingFace models | none needed | +| `remote` | `registerModel({ modelId, model, modelLib })`, any base URL | none — a bad URL is reported far better by WebLLM's loader than by a HEAD request | +| `injected` | `registerModel({ modelId, files })`, **no network at any point** | exhaustive, before the first byte is written | + +Three things fell out that were not obvious going in: + +1. **It cost almost no code.** `toAppConfig` already emitted `{model, model_id, model_lib}`; it had + only ever been handed `.invalid` URLs. +2. **It retired the top risk.** Cache injection on an ordinary page origin was risk #1 because the + whole design rested on it. It is now the offline route only. +3. **The routes cannot be confused, so downloads default on.** `.invalid` is reserved by RFC 6761 + and can never resolve, so an injected model whose cache was evicted fails with a DNS error rather + than silently pulling a gigabyte. That is the *mechanism* of the offline guarantee, not a label + for it — a test asserts every URL such a record carries is on a `.invalid` host. + +**Why local upload keeps a synthetic origin.** WebLLM composes every artifact URL as +`new URL(relative, base)` and runs the base through `cleanModelUrl`, which calls `new URL(...)` — so +the base must be absolute and resolvable. A `blob:` URL cannot serve as one, and there is no hook to +hand the loader bytes directly. Pre-populating the cache under WebLLM's own scopes and keys *is* its +native path. Seeding the cache for a URL the developer hosts was rejected: it would unify the record +shapes but make eviction silently re-download a gigabyte, which is the failure the design prevents. + +## Model lifecycle — four states, three operations + +The old API collapsed them, and that was a real bug: `remove()` deleted the cache **and** the +registry record, so for a remote model it threw away the only URL the bytes could be fetched from. + +| state | VRAM | cache | record | leave via | +| --- | --- | --- | --- | --- | +| resident | ✅ | ✅ | ✅ | `unload(id)` | +| cached | — | ✅ | ✅ | `evict(id)` | +| registered | — | — | ✅ | `remove(id)` | +| unknown | — | — | — | — | + +**Multiple resident models.** `#pool` became `#pools: Map` with a current +selection. `use(id)` switches for free; a request naming a resident model routes to it *without* +changing which is current. Additive residency is opt-in (`{ keepResident: true }`) because each +resident model is a full copy of its weights and nothing reports free VRAM to a page. + +## Device and compatibility + +`probeDevice()` / `canRun()` / `rankModels()` answer "will this run here" before a byte is fetched. +The rules are this project's platform scars as code: the blocklisted adapter, Firefox's +9-storage-buffer cap, `q4f16_1` on a device without `shader-f16`. + +Two things learned while building it: + +- **Blockers and warnings must stay separate.** The 9-buffer cap costs KV reuse but is a *warning* — + blocking it would refuse the exact configuration this project ships on. +- **"Largest that fits" is bad default advice.** Decode is memory-bandwidth-bound, so the largest + model that fits is also the slowest thing that fits. `prefer: "quality" | "speed"` makes it the + caller's choice rather than an assumption. + +`probeDevice()` never throws — an unusable device is a result to explain, not an exception. + +## Stop reinventing WebLLM + +An audit ([WEBLLM-SURFACE.md](WEBLLM-SURFACE.md)) found three functions reimplemented that WebLLM +exported all along. Cause: the bundle was treated as something to `grep` for narrow facts rather than +an API to survey once — `cleanModelUrl` was even *read on screen* and then rewritten, without asking +what else used it. This violated AI.md's own **Reuse First** principle. + +| Removed | Replaced by | +| --- | --- | +| `cleanModelUrl()` reimplemented | nothing — it only propped up the two below | +| `ModelStore.cacheKeysFor()` | WebLLM derives the keys it fetched | +| `ModelStore.cacheState()` | `engine.cacheState()` → `hasModelInCache` for remote | +| `ModelStore.evict()` (all sources) | `evictInjected()` + `engine.evict()` → `deleteModelAllInfoInCache` | +| speed re-derived from the worker probe | `usage.extra.decode_tokens_per_s`, already on every response | + +The routing rule is now explicit: **whoever wrote the bytes owns the keys.** Our path survives only +where it demonstrably does more — WebLLM's delete and cache-check both read `tensor-cache.json` to +enumerate shards, so once *that* file is evicted they can neither find nor clean the shards it +indexes. Injected records carry an explicit key list and have no such failure. There is a test for +exactly that case, and it is the only justification for keeping the code. + +**Speed was worse than duplication:** the pool already set `include_usage` and already stored +`chunk.usage`, so the measurement was being *received and discarded* so the worker probe could +recompute it. + +## Raw chunk pass-through, and the tool-calling bug + +The pool stripped every chunk to `delta.content` and the facade rebuilt an envelope from scratch — +so `tool_calls` was dropped entirely (**tool calling returned nothing usable**), `logprobs` was +always `null`, and `created` was restamped per chunk. Chunks now pass through verbatim. + +Nothing is synthesized on the normal path: WebLLM emits its own terminal `finish_reason` chunk and +its own usage chunk. The one exception is an interrupted generation, where the stream simply stops +and a consumer would otherwise never learn why. + +**A correction to the plan that produced this.** It specified `mergeToolCallDeltas()` for +"standard OpenAI fragment accumulation". WebLLM does not stream fragments — it parses the whole +output message at the end and emits tool calls complete in one terminal chunk. Building the merge +would have been machinery for a wire shape that is never produced: the plan's own failure mode, +inside the plan meant to prevent it. + +## Bugs found and fixed along the way + +| | | +| --- | --- | +| `EnginePool.load()` leaked an engine | It awaited `createEngine` before assigning `#slots`, so an `unload()` in that window tore down an *empty* pool and the engine then installed itself into a pool nobody referenced — leaking a worker and a full copy of the weights. `#grow()` had always guarded this; `load()` never did. Fixed with a `#generation` counter. | +| `registerModel` accepted URLs that fail at load | WebLLM's `cleanModelUrl` ends in `new URL()` with no base, so a relative `/models/x/` throws deep in the loader. Now resolved at registration. | +| `state.modelId` / `resident` went stale | Views onto `#pools` that nothing re-synced after unload, and left pointing at a model that never came up after a failed load. | +| `probeDevice` threw on a partial `navigator.gpu` | A polyfill without `requestAdapter` produced a TypeError from a function documented never to throw. | +| The pool discarded `finish_reason` | A `max_tokens` truncation was indistinguishable from the model choosing to stop. | +| `store.remove()` stranded remote shards | It iterates `groupKeysByScope`, empty for remote records — and deleting the entry destroys the only URL those bytes could be derived from. Now `engine.remove()` evicts first. | +| `throw`-as-`goto` in `multistep.js` | Caught two lines below; replaced with the control flow it was emulating. | +| `features()` called a dead fast path healthy | It answered "is multi-step on?" with `decodeSteps > 1`, but the worker keeps posting the *configured* K — 15 — long after the contract check routed decoding to stock single-step. So the one call documented as "what is switched **on** now" reported `multiStepDecoding: true` for the exact fault it exists to surface, and never exposed `multiStepOff` at all. `environment()` escaped it only by reading `state.decode.multiStepOff` itself rather than trusting `features()`. Now `multiStepDecoding` is `decodeSteps > 1 && !multiStepOff?.length`, and `multiStepOff` is returned beside it. `decodeSteps` deliberately keeps reporting the knob's value: dropping it to 1 would make `environment()` advise `configure({ decodeSteps: 15 })` for a fault no setting can fix. | + +## Corrections to the record + +Kept because each was stated confidently and was wrong; a future reader should not re-derive them. + +- **"A second engine measured 1.06x, so batched decode is the only route to concurrent throughput."** + This framed two complementary mechanisms as substitutes. A second engine buys *task isolation* — + a translation and a ghost-text completion running at once — and never claimed aggregate + throughput; the GPU is already saturated by a 2B. Batched decode makes *one task's* many requests + faster. Neither substitutes for the other. +- **"Multi-model residency via `#pools` was not necessary."** Wrong. `reload()` unconditionally calls + `unload()` first, so `reload([a,b])` is all-or-nothing — adding a third model reloads the first two + (~51 s each). Additive residency does not exist upstream. +- **"No load time is measured in the repo."** It is: 51 s, [AI.md](AI.md) line 77. A grep for the + wrong phrasing missed the table row. + + +## `npm run e2e` verified the extraction on real hardware + +First run against the post-extraction tree: real Firefox, real GPU (Apple Silicon, `shader-f16`), +real `Qwen3.5-0.8B-q4f16_1-MLC`, drag-and-drop ingestion through the production `src/engine/` and +`src/adapters/webext.js` paths. **`e2e PASS`.** + +- Ingest: 443,129,354 bytes, 11 shards, 2531 ms. Load: 48.1 s (AI.md's 51 s figure is for the larger + 2B; the 0.8B here loading faster is consistent with that being memory-bandwidth-bound). +- Decode: 27.4 tok/s over 127 tokens — inside AI.md's measured 16.6–27.9 tok/s range for this model. + Decode probe: 664 kernels/token (639 forward + 25 sampling), 16.1 flushes/token — the same shape + the compute-pass-batching patch targets, and it is still applying (41.3 kernels/flush). +- KV-reuse path exercised and correct: paged prefill and forced ragged re-prefill produced identical + output on a multi-round conversation. Re-prefill slope 2.29 ms/token, in the neighbourhood of + AI.md's 5.27 ms/token figure (different model, different history length — not a direct comparison). +- The scheduler's own two-tasks-two-engines check: 3.8 s concurrent vs 4.0 s sequential = **1.05x**, + consistent with AI.md's measured 1.06x. This is the number the "second engine buys isolation, not + throughput" framing rests on, now reconfirmed after the pool moved to `#pools: Map` — evidence the multi-model split did not regress the single-model scheduling behaviour + it was built on top of. + +**One number worth a second look, not treated as a finding here:** this run reported +`storageBuffersPerStage=9`. It did not block anything — the KV-reuse path was exercised in the same +run and passed — but it is the exact threshold `probeDevice()`'s `NO_KV_REUSE` warning keys off, so a +future run reporting the same value is worth cross-checking against `device.test.mjs`'s assumptions +rather than assumed benign a second time. + +> **[resolved] A second run reported 9 again, and the cross-check says 9 is the baseline, not an +> anomaly.** AI.md has said so all along — the Firefox Metal backend caps +> `maxStorageBuffersPerShaderStage` at 9, which is the entire reason the `storage-buffer-limit` +> patch exists. What the cross-check *did* surface is sharper than the original worry: +> `device.test.mjs` defines a healthy device as `storageBuffers = 10`, so the "good" default in the +> test matrix describes hardware nobody here has. On the reference M4, `probe.kvReuse` is +> **always** `false`, `engine-worker.js` forces `resetChat()` on every prefill, and +> **`batch_prefill_paged_kv_kernel` has therefore never executed on real hardware** — it is +> mock-tested only. It would first run on a >=10-buffer device, i.e. Chrome (ROADMAP, Gate B). +> +> That also made the e2e's own multiround check misleading: with reuse forced off, its "with KV +> reuse" and "forced reprefill" branches both ran the ragged kernel, so `identical` was guaranteed +> and the line `paged prefill is fine` claimed something the run had not tested. The check now +> reports `UNVERIFIED for paged prefill` on a sub-10-buffer device and still fails if two ragged +> re-prefills of the same history disagree. + +The manifest.json restore left a diff — `restore()` round-trips the file through +`JSON.parse`/`stringify`, which turns `\uXXXX`-escaped em-dashes back into literal UTF-8. Cosmetic, +not a behaviour change, reverted with `git checkout`. Worth knowing before the next e2e run leaves the +same diff and it looks like something broke. + +> **[fixed]** It did leave the same diff on the next run. The snapshot now only round-trips through +> JSON when the tree is *actually* dirty (a run killed mid-flight leaves the patch behind); a clean +> file is restored byte-for-byte. The cost of the old behaviour was not untidiness — it was training +> the reader to ignore a dirty tree after an e2e, which is precisely when a real diff matters. + + +## Lossless WebLLM upgrade — a bump is minutes, not an afternoon + +`@mlc-ai/web-llm` is pinned exactly because `build.mjs` rewrites the bundle by matching source text. +"Lossless" was never meant as "automatic" — it means a bump *fails at the right line* instead of +somewhere deep in a half-patched loader. The standing runbook is in +[WEBLLM-SURFACE.md](WEBLLM-SURFACE.md), "Upgrading"; this is why each piece exists. + +Three guards, because there are three distinct ways an upgrade breaks us: + +| drift | caught by | the failure it prevents | +| --- | --- | --- | +| **surface** — code moved or reformatted | `build/patches.mjs` verify-then-write | a patch anchor silently landing in the wrong place, or the build half-applying and reporting only the first miss | +| **semantic** — a symbol survives, its meaning changed | `test/webllm-contract.test.mjs` | an export deleted, a field renamed, an enum gaining a case — none of which throw | +| **behavioural** — every name and shape intact, output wrong | `npm run e2e` | a tvmjs refactor that changes numerics | + +**Contract tests (2a).** Makes WEBLLM-SURFACE.md executable — every export and shape the project +depends on, asserted statically against the bundle, GPU-free, first in `npm test`. The +highest-value piece: it catches semantic drift, which the patches cannot see. Two guards beyond the +obvious list: the monkeypatch member list is *derived from `multistep.js`'s own source* so it cannot +go stale, and `model_lib` unguessability is asserted rather than assumed (if it became derivable, +the "do not guess" rule in the verb-consolidation section should be revisited). Each assertion class +was mutation-tested — which found a real bug: `bundle.includes(name)` still passes when +`processNextToken` becomes `processNextTokenV2`, since the old name stays a substring. Now +word-bounded. + +**Patch self-check and fuzzy diagnostics (2b).** Patches moved to `build/patches.mjs` as data, +applied by a shared verifier. Every anchor is checked before anything is rewritten. `patch-manifest.json` +records the version the anchors last held against, so a bump announces `0.2.84 -> 0.2.85: verifying +4 anchor(s)` rather than silently succeeding. Vanished identifiers are matched against survivors by +trigram overlap — simulated upstream renaming `requiredMaxStorageBuffersPerShaderStage`, the +diagnostic found the replacement at 85% similarity with its line. Ambiguity is a hard stop too: an +anchor matching 1995 sites refuses rather than rewriting one at random. Two corrections the +simulation forced: rank candidate lines by summed *rarity* not hit count (raw count returned +`const msg = {` — true and useless), and a rename needs a human to *approve* the new anchor, not to +*find* it. + +**Structured patches (2c).** The corrected expectation held: AST parsing survives *formatting* +drift, not renames — an AST search by name fails exactly as a string match does. So the gain is +narrower than "structured = durable", and the work matched the correction rather than the original +proposal. `in: { enclosing }` scopes an anchor to the function a *sibling anchor* matched in — a +matched anchor, not a function name, so it adds no identifier upstream could rename. This removes the +false-failure class around `compute.end();`, a string generic enough that any unrelated new compute +pass in tvmjs failed the build. Anchors also match modulo whitespace and are word-bounded — the +latter not in the plan and found the same way 2a's bug was: `compute.end();` is a substring of +`precompute.end();`. The rebuilt bundle is byte-identical to the string-replacing applier's output. +`acorn` is a devDependency, ~565 KB unpacked (the original estimate was off 10x), never shipped. + +**Runtime monkeypatch guard (2d).** `multistep.js` drives ~30 undocumented tvmjs internals; a rename +turns the fast path off *silently* — ~18 -> ~10 tok/s with nothing in the log. 2a covers the static +half. The runtime half: `PIPELINE_CONTRACT` + `missingPipelineMembers()`, checked against the live +pipeline at first decode (there is no pipeline at install time — the engine gets one per `reload()`), +verdict cached. Three buckets, because presence is not the failure that hurts: `calls` must be +callable (a rename throws — loud), `numbers` must be numbers (`x += 1` on an absent member creates a +property and the KV accounting drifts — silent), `reads` need only exist. Failure is loud once per +pipeline and posts `multiStepOff` to the host, which is otherwise indistinguishable from an idle +engine since `onBurst` is the only thing that reports stats. Found while building it: 2a's derivation +matched `\bpipeline\.` and missed members reached across a line break — the softmax at the heart of +the burst had no rename guard for as long as that test existed. Now whitespace-tolerant, comments +stripped first, and `PIPELINE_CONTRACT` is asserted to equal what the source actually reaches for. + +**The flow, documented (2e).** Moved to WEBLLM-SURFACE.md so the doc you must revise on a bump is +the doc that tells you how. + + +## Verb consolidation — the ergonomic layer + +`chat.completions.create()` is the compatibility layer and never changes. Everything here is +*additional* — the verbs a developer reaches for when they are not porting WebLLM code. + +**`load(src, opts)` — one polymorphic entry.** Absorbs `load` + `registerModel` + +`ingestModelFolder`. Dispatch is a pure, synchronous `classifySource()` in +[src/engine/sources.js](src/engine/sources.js), so all four shapes — prebuilt id, HF/hosted URL, +`{model, modelLib}`, folder off disk — are testable with no GPU and no store. `registerModel` and +`ingestModelFolder` stay exported unchanged; `load()` composes them. + +Two dispatch rules were **dropped after measuring**, both because a wrong guess surfaces as a 404 +deep inside the loader: + +- **`modelLib` is never guessed.** `-webgpu.wasm` matches **0 of 163** prebuilt models + (real names carry a `_cs1k`-style suffix, drop `-MLC`) and **0 of 163** host the lib on the + weights' origin (they live on `raw.githubusercontent.com`). A remote source without `modelLib` + fails in the classifier with that sentence, before any fetch. +- **`/resolve/main/` is not derived for HF URLs.** WebLLM's `cleanModelUrl` already appends it; doing + it ourselves double-applies. A test asserts the stored URL is byte-identical to what was passed. + +The id *is* derived from the URL's last segment — safe where `modelLib` is not, because an id is a +key in our own registry, never a path anything fetches, so a wrong guess is visible immediately and +free. `{ id }` overrides. `defer: true` on a bare prebuilt id is an **error**, not a silent load — +that silent load is the trap the whole section exists to avoid. Unknown ids get near-match hints. + +Six mutation tests. One false pass worth remembering: the near-match hint appears at **two** error +sites and `String.replace` mutated only the first, so a working guard looked untested — a mutation +that does not apply is indistinguishable from a guard that does not work. Two latent crashes fixed +on the way: `filesFromInput`/`filesFromDataTransfer` spread their argument, so an array-like-but-not- +iterable `FileList`/`DataTransferItemList` died with `fileList is not iterable` three frames from +the caller's drop handler. Both use `Array.from` now. + +**`unload(id, level)` — two depths, not two verbs.** `UNLOAD_LEVEL` is `"vram"` (default: free +VRAM, keep cache + record) or `"cache"` (also delete bytes, keep record — the old `evict()`). +`remove()` keeps its own verb: it is the one that cannot be undone without re-supplying the source. +An unrecognised level is refused with an error pointing at `remove()`, because "forget this model" +is the reading someone will try to spell as a level and it is the destructive one. + +**[settled] A bare `unload()` frees only the current model**, with `unloadAll()` explicit for the +rest. The plan had the bare call free *everything*; shipping that silently would trap anyone already +calling `unload()`. "Free everything" is the more destructive reading and should be asked for by +name. `#evictBytes()` was split out of `evict()` so `unload(id, "cache")` reaches the bytes without +re-entering the class for a pool just torn down. + +**`environment()` — read-only report; `.measure()` on it.** Absorbs `probe` + `features` + +`estimateSpeed` for the *reporting* half. [src/engine/environment.js](src/engine/environment.js), a +callable `engine.environment` cached like `chat`. Three open questions were all resolved by one +decision — **split read from write**: `environment()` reports only, writes go through `configure()`, +and passing a setting to `environment()` is an *error naming `configure()`*, not a silent no-op. +Implicit read/write dispatch by argument shape is the opposite of foolproof — the "reject or +write-then-report?" question had no intuitive answer precisely because one function was doing two +jobs. + +Every report line carries `severity` · `affects` · `cause` · `fix` · `operable`, with **`fix: null` +⟺ `operable: false`** asserted for every line — hardware, build-time flags and browser settings +report a consequence with no remedy, which is still the difference between a bug report and an +informed decision. A **blocked device short-circuits the report**: "K=15 forward steps per GPU sync" +next to "no model can load" is true and useless. `configure()` grew `engineCount` because the report +advertises it as operable and a report naming a call that throws is worse than no report — it is +persisted, not hot, and `environment()` reports that gap rather than pretending. The `multiStepOff` +guard (§2d) finally has a consumer: a `degraded` line naming the missing internal, where before it +was posted by the worker and read by nothing. + +Seven mutation tests. One real hole found: "`local` never fetches" was tested with a fetch counter, +but `load()` caches the model's size so `estimateSpeed()` short-circuits and *neither* scope fetches +after a load. The guarantee is structural — `local` never consults the model layer — and is tested +that way now. + + +## Engine capability — prefetch, embeddings, recipes + +**`prefetch(modelId)`** — [src/engine/prefetch.js](src/engine/prefetch.js). Downloads a model with +no engine and **no WebGPU at all**: an app can warm the cache before it knows whether the machine +can run the model. Resumes; a second call is free. + +The hard part: fetching the artifacts ourselves means deriving their URLs — the `/resolve/main/` +rule the verb-consolidation work above refused to derive. That refusal still holds; it was about not +deriving a URL *WebLLM will derive again at load*, which double-applies. Here WebLLM is not in the +loop — we are the loader. What makes it safe is not trusting the derivation: a key off by one +character writes a cache the loader never reads, and prefetch would report success while the user +downloads the model twice. So every prefetch ends by asking **WebLLM's own `hasModelInCache`** — +which derives through the very function we mirror — and throws if it says no. The contract test also +pulls `cleanModelUrl` out of the bundle and *runs* it against ours on six URL shapes, so an upstream +scheme change fails a test, not a download. Seven mutation tests, all caught. + +**Embeddings (`engine.embed()`)** — a `kind` field on the job and one branch in +[pool.js](src/engine/pool.js) `#start`. **One pool, not two**: priority, supersession, preemption +and one-task-one-engine are identical for both kinds; only the call at the far end differs. A second +pool would have duplicated the scheduler to change one line. `embed()` returns bare vectors, +`embedRaw()` keeps WebLLM's envelope. **Known limit:** a running embedding cannot be interrupted — +`interruptGenerate()` works by making a decode loop break out, and one forward pass has no loop, so +a cancel that lands after the job starts marks it cancelled without stopping it. Stated in the +JSDoc, the README and a `[known limit]` test rather than left to be discovered. Six mutations, five +caught; the sixth was *equivalent* — `#start` decides on an explicit `=== EMBEDDING` and any unknown +kind routes to chat either way. + +**Recipes — `ask()`, `conversation()`, `ghostText()`** — [src/engine/recipes.js](src/engine/recipes.js), +also methods on the engine. Scope grew on request: one command for each of the three things apps +actually want. The scheduling shipped as specified — one stable `session`, `interactive`, short +`max_tokens`, debounce, `cancel()` on blur, stale contexts dropped — and **prompts stayed with the +caller**: `ghostText({ prompt })` is required with no default; `ask`/`conversation` carry the +caller's text through. The engine authors nothing. + +The piece worth keeping: `suggest()` **resolves `null` when stale**. The engine already superseded +the work; what a caller still had to remember was not to *paint* the answer that came back anyway. +Returning `null` removes the choice — the difference between a policy and a wrapper. +`conversation()` bounds history at 12 exchanges, derived from AI.md's numbers: with no cross-turn KV +reuse every turn re-prefills at ~5.27 ms/token, so unbounded history is quadratic and a turn near +the limit waits ~22 s. `keep: Infinity` opts out. + +Found and fixed: a **promise leak in the debounce**. A newer keystroke called `clearTimeout` on the +previous waiter, whose `await` then had nothing to resolve it — every superseded keystroke leaked a +promise that never settled, and `Promise.all` over a burst hung forever. A superseded waiter has to +be woken and told it lost, not merely disarmed. Twelve mutation tests; two initially passed — one +equivalent, one genuinely vacuous: `sent[0].session === sent[1].session` also holds when *neither* +has a session, which is exactly the regression it was meant to catch. Presence is asserted before +equality now. + + +## Shipping 0.1.0 — installable, documented, on npm + +The library was extracted, tested and complete, and served the project's goal — "make WebLLM easier +to use, foolproof to build on" — **for nobody**, because it was unpublished, undocumented as a whole +surface, and un-installable. This section is the gap between "the code is done" and "a developer can +`npm i` it and run four lines." + +**Four-line target, met without an engine change.** `import` / `CreateScheduledEngine(id)` / +`engine.ask(prompt)` / read the string. Probing the shape found three things that stopped it being +*usable*: + +1. **Un-installable.** `vendor/web-llm.js` is a build product and is gitignored; there was only a + `prepublishOnly`, and npm runs **`prepare`** for a git dependency. So `npm i` 404'd (`"private": + true`) and a git dependency installed but could not run, failing with `GENERATION_FAILED: Cannot + find module .../vendor/web-llm.js` — wrong twice, since nothing had begun generating and the path + named is ours. Now `"prepare": "node build.mjs"`, `private` removed. Verified by deleting + `vendor/` and running `npm install` — it comes back. +2. **Looked like a hung process.** No `initProgressCallback` meant zero output during a minutes-long + ~0.8 GB download. `CreateScheduledEngine` now distinguishes three states: `undefined` → a + throttled console reporter (1 line/second; 58 shard callbacks → 2 lines; the 100% report is never + dropped), `null` → explicit silence, a function → unchanged. **`new ScheduledEngine()` stays + silent** — a library core that logs is wrong in a worker, an extension background page, or a + test. This is the getting-started facade only. +3. **The Vite worker-URL break.** Vite's dependency pre-bundler copies `new Worker(new + URL("./engine-worker.js", import.meta.url))` verbatim into `node_modules/.vite/deps/`, where the + sibling file does not exist — `vite dev` only, real (non-linked) install only, `vite build` + unaffected. `everything-webgpu/vite` ships a plugin (`optimizeDeps.exclude`, the manual + equivalent still documented). And if a consumer does neither, `load()` now fails with + `PACKAGE_INCOMPLETE` naming the fix, because `new Worker()` does not throw on a 404 — it fires one + `error` event and goes quiet, so the handshake is raced against it. + +**`PACKAGE_INCOMPLETE` is one code for two causes** (`detail.cause` separates them). No caller +writes a different `catch` branch: both mean "your build is wrong, this app has not shipped," both +are fixed in config. A second code would grow the table a caller reads without giving them anything +to do. + +**`verify-consumer` — the only test that can see the consumer's world.** Everything under `test/` +and every `examples/` project reaches the package through a *linked* path, and Vite never +pre-bundles a linked package — so none of them can exercise the one failure that reaches users. This +blind spot produced a **wrong claim in the docs**: that `optimizeDeps.exclude` was needed for `vite +build` and that the examples proved it. Measured on a real tarball install, neither holds — build +output is byte-identical with and without it. `npm run verify-consumer` packs the tarball, installs +it for real, and asserts three outcomes separately: `vite build` emits the worker chunk and keeps +WebLLM lazy; `vite dev` **without** the plugin still 404s the worker; `vite dev` with it resolves. +The middle one is asserted as a *failure* on purpose — a fix whose absence changes nothing is not a +fix, and if Vite ever stops pre-bundling this package that assertion says the plugin is dead weight. + +**`API.md` — every call form on one page,** asserted by +[test/api-doc.test.mjs](test/api-doc.test.mjs), derived from the source the way `readme.test.mjs` +is: every `engine.x(` named resolves to a real member, **no public member is left undocumented** +(the reverse direction the README test lacks), the error table equals `ERROR`, every export +appears, enum-value rows match the real objects, the subpath table equals `package.json` `exports`. +Writing it found `engine.store` undocumented and a regex reading enum rows as error codes. + +**`examples/` — bare, react, webext,** each a standalone project depending on the package as +`file:../..` so it resolves through the **`exports` map** — an example importing +`../../src/engine/index.js` would still run and would still leave the exports map, the `files` list +and every entry path untested. [test/examples.test.mjs](test/examples.test.mjs) derives its checks +from the example sources, so a fourth example is covered the moment its directory exists. Also +closed a silent `files`/`exports` gap: a new export path that `files` would not publish resolves in +the checkout and 404s in the tarball. Asserted from `package.json` now. + +**Bundle-size story, measured not estimated:** **53 kB (~19 kB gzip)** entry chunk before a model +loads; the 6 MB WebLLM bundle is a **lazy chunk** fetched on the first `load()` or +`listAvailableModels()` and never by a visitor who does neither; the IndexedDB adapter is a further +0.8 kB lazy chunk that vanishes when a host brings its own store — the webext build emits no `idb` +chunk at all, which is that claim tested by construction. + +**Licence compliance.** Publishing `vendor/web-llm.js` redistributes WebLLM (Apache-2.0) and its +dependency `loglevel` (MIT), and the esbuild bundle was built `legalComments: "none"` — no notice +survived, a violation. `THIRD-PARTY-NOTICES.md` now carries the full texts, generated from the +installed packages; [test/license.test.mjs](test/license.test.mjs) fails the build if a bundled +dependency ever lacks a notice, catching a future web-llm bump that inlines a new dep. `build.mjs` +uses `legalComments: "eof"` now — upstream has already stripped every `@license` banner (the bundle +is byte-identical either way today), but `"none"` would silently drop one a future dep adds. `LICENSE` +added (ISC). `files` scopes `vendor` to `web-llm.js` — the stale, unreferenced `vendor/web-llm.d.ts` +was shipping and made the tarball depend on disk state. + +**Published:** `everything-webgpu@0.1.0`, 46 files, 2.3 MB packed; `dist.integrity` matched the +dry-run exactly. Still deferred to a later version: the `demo` extension rebuilding on the package +(the source-tree acceptance test for the extraction), and a `repository` field once the repo has a +remote. diff --git a/LICENSE b/LICENSE new file mode 100644 index 0000000000000000000000000000000000000000..873032b51194f072538d3cac9a51d535ca3b96fd --- /dev/null +++ b/LICENSE @@ -0,0 +1,21 @@ +ISC License + +Copyright (c) 2026 nyaaorick + +Permission to use, copy, modify, and/or distribute this software for any +purpose with or without fee is hereby granted, provided that the above +copyright notice and this permission notice appear in all copies. + +THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES +WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF +MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR +ANY SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES +WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN +ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF +OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE. + +--- + +This distribution includes a bundled build of WebLLM (@mlc-ai/web-llm, +Apache-2.0) at vendor/web-llm.js and its transitive dependency loglevel +(MIT). Their license texts are in THIRD-PARTY-NOTICES.md. diff --git a/MLC-COMPILE.md b/MLC-COMPILE.md new file mode 100644 index 0000000000000000000000000000000000000000..97ce4e19a9239818d2bfcd1ec5d3c61a77642b05 --- /dev/null +++ b/MLC-COMPILE.md @@ -0,0 +1,292 @@ +# Qwen3.8-2B-Distill → MLC/WebGPU: built, and what it cost + +**Status: done.** `dist/Qwen3.8-2B-q4f16_1-MLC/` loads in Firefox and passes `npm run e2e`. +Everything labelled **measured** below was verified on this machine (M4 MacBook Air, 16 GB, +Firefox 154, macOS 26.5.2) during the session that built it. + +The compile itself was the easy part. The toolchain was not: **no published pair of MLC nightlies +works together**, and the fix is not a version pin. See "The toolchain does not ship working" — +read it before touching the venv, because a fresh `pip install` reintroduces every one of these. + +## What was produced + +| | | +| --- | --- | +| weights | 26 shards, **1.059 GB** on disk (1,089,449,808 bytes materialized) | +| params | **1,881,825,088** at **4.503 bits/param** (0.987 GB post-quantization) | +| library | `Qwen3.8-2B-q4f16_1_cs1k-webgpu.wasm`, 6,889,261 bytes, 154 kernels | +| config | q4f16_1, context 4096, prefill chunk 1024, conv template `qwen3_5` | + +Reproduce with `tools/setup-mlc-toolchain.sh`, then: + +```sh +python3 tools/strip-vision.py --keep-prefix # 632 -> 320 tensors, 1.88B params +.venv-mlc/bin/python -m mlc_llm gen_config model/Qwen3.8-2B-text \ + --quantization q4f16_1 --conv-template qwen3_5 \ + --context-window-size 4096 --prefill-chunk-size 1024 -o dist/Qwen3.8-2B-q4f16_1-MLC +.venv-mlc/bin/python -m mlc_llm convert_weight model/Qwen3.8-2B-text \ + --quantization q4f16_1 --device cpu -o dist/Qwen3.8-2B-q4f16_1-MLC +. ~/emsdk/emsdk_env.sh +export TVM_HOME=~/.cache/mlc-web-build/tvm +export MLC_LLM_SOURCE_DIR="$PWD/.venv-mlc/lib/python3.12/site-packages/mlc_llm" +.venv-mlc/bin/python -m mlc_llm compile dist/Qwen3.8-2B-q4f16_1-MLC/mlc-chat-config.json \ + --device webgpu -o dist/Qwen3.8-2B-q4f16_1-MLC/Qwen3.8-2B-q4f16_1_cs1k-webgpu.wasm +``` + +`--device cpu` on `convert_weight` is not a preference. **Measured:** `--device metal` dies in +TVM's LLVM ORC JIT with `Unexpected definitions in module TVMMod-jitted-objectbuffer` on the first +quantize kernel. The CPU path quantizes all 1.88 B params in 7.4 s, so there is nothing to miss. + +## Measured performance + +``` +decode 16.9 tok/s over 127 tokens (projection was 14-18 — held) +prefill 47.9 tok/s, ttft 0.71 s +load ready in 51.1 s (ingest 4.8 s of that) +decode probe 664 kernels/tok = 639 forward + 25 sampling + cpu-encode 3.9 ms/tok, gpu+tick 62.7 ms/tok + 16.1 flushes/tok -> 41.3 kernels per flush +batch 4 prompts, 96 tok in 9.7 s = 9.9 tok/s aggregate at pool size 1 +``` + +`639 forward kernels ÷ 24 layers = 26.6`, i.e. the same ~27 the 0.8B showed. dlight made the same +call here — more reduction threads over more work per thread — so the previously measured **1.83x +from going 2 -> 32 iters/thread is still on the table**, and the toolchain is now up to try it. + +### The cost of having no KV reuse — now measured + +This was the handoff's one explicitly-unmeasured risk ("measure it before assuming it is fine for +long chats"). It is not fine for long chats: + +| history | wall | tok/s | +| --- | --- | --- | +| 29 tok | 624 ms | 71 | +| 249 tok | 2,604 ms | 104 | +| 915 tok | 9,251 ms | 101 | +| 2,249 tok | 12,318 ms | 186 | + +**Slope: 5.27 ms per history token (190 tok/s marginal), floor 624 ms.** Every turn re-prefills the +whole conversation, so at the 4096-token context limit a late turn pays roughly `624 ms + 5.27 ms × +4000 ≈ 22 s` before its first token. Short chats are unaffected; long ones are not usable. If this +model is meant for long conversations, restoring KV reuse (see below) is no longer optional. + +## The storage-buffer goal: what actually happened + +The aim was a build where nothing the engine dispatches binds more than 9 storage buffers, reached +**by deleting** dead kernels. Deletion turned out not to be available. + +**Measured** — the 2B's histogram is identical to the verified 0.8B's, kernel for kernel: + +| storage bindings | kernels | +| --- | --- | +| 1-6 | 146 | +| 8 | 1 | +| 9 | 3 | +| **10** | **4** | + +The same four at 10: `batch_prefill_paged_kv_kernel`, `batch_prefill_paged_kv_sliding_window_kernel`, +`batch_tree_attn_kernel`, `tree_attn_paged_kv_kernel`. + +**There is no flag, and there cannot be one.** In +`tvm/relax/frontend/nn/llm/kv_cache.py`, `_create_tir_paged_kv_cache` registers the sliding-window, +tree-attention and paged-prefill kernels unconditionally, as a **fixed positional tuple** handed to +the C++ `PagedKVCache` constructor. The runtime indexes that tuple by position, so omitting an entry +does not delete a kernel — it misaligns every kernel after it. This is the case the original handoff +anticipated ("if MLC emits the full set unconditionally, that is acceptable"), now confirmed at the +source rather than assumed. + +So the position is unchanged from the 0.8B: three of the four are dead by config, and the live one, +`batch_prefill_paged_kv_kernel`, stays unreachable only because of the engine-side guard in +`src/background/engine-worker.js` (`resetChat()` before every prefill when +`maxStorageBuffersPerShaderStage < 10`). `tools/audit-wasm.mjs` still exits 1, correctly — it judges +the file, and it cannot see a guard that lives in the engine. + +**Measured, and worth stating precisely:** the e2e's multi-round A/B came back *identical*. That +proves there is no silent garbage. It does **not** prove paged prefill works — the guard means both +branches of the A/B ran through the ragged kernel. The harness's own "paged prefill is fine" wording +overstates it. + +### If you want KV reuse back + +Unchanged from before, and now motivated by the measured 5.27 ms/token slope. Do not re-fuse +anything: `batch_prefill_paged_kv_kernel` binds six `array` metadata buffers whose uniform block +already carries `*_elem_offset` for each, and the body already indexes as `name[expr + name_elem_offset]`. +Packing them into one buffer with six offsets is a signature change, not an algorithmic one: 10 -> 5. +Merging any two reaches 9. Grep `q_indptr` / `page_values` / `k_rope_pos_offset` in TVM — with the +toolchain now standing up, this is a tractable change rather than a research project. + +## The toolchain does not ship working + +Five independent breakages, none of which is a version-pin problem. `tools/setup-mlc-toolchain.sh` +applies all of them; `tools/patch-mlc-nightly.py` and `tools/resign-venv.sh` are idempotent and must +be re-run after **any** `pip install` touching these packages. + +**1. Every wheel dylib has an invalid signature.** `import tvm` is SIGKILLed by the kernel +(`CODESIGNING / Invalid Page`) inside `dlopen`. The ad-hoc signatures do not match the bytes. +`codesign --force --sign -` over each recomputes them. Silent — no Python traceback, just exit 137. + +**2. `apache-tvm-ffi` must be pinned to exactly `0.1.13`.** The metadata says `>=0.1.13`, which +resolves to `0.1.13.post3`, whose ABI does not match the bundled runtime: it corrupts the heap in a +static initializer and aborts. (`0.1.13` is marked yanked on PyPI. It is still the one that works.) + +**3. `psutil` is missing from the wheel's dependencies** but imported by `mlc_llm.serve`. + +**4. The published nightlies are mutually incompatible — this is the big one.** +`mlc-llm-nightly-cpu 0.26.dev6` is written against an *unreleased* TVM in which `PrimFunc.params` +holds `Buffer`s directly and `tirx.is_buffer_var` exists. Every published `mlc-ai` wheel +(dev61/dev203/dev246) still has `params: Array` plus a separate `buffer_map` — verified in +`include/tvm/tirx/function.h` — and none defines `is_buffer_var`. Going the other way fails too: +`mlc-llm 0.20.dev162`'s Python *does* use `buffer_map`, but its dylib needs `TVMBackendParallelLaunch`, +which dev246 no longer exports. + +dev6 is kept, because it has the newest `qwen3_5` model definition and the correct `qwen3_5` +conversation template. The places where it reaches for the unreleased API are rewritten against +`buffer_map` — **not invented**: each rewrite restores what `mlc-llm 0.20.dev162` does in that same +function, kept on dev6's relax spelling (`TensorType`/`ty_args`, where dev162 said +`TensorStructInfo`/`sinfo_args`). Affects `low_batch_specialization.py` and +`lift_global_buffer_alloc.py`. + +Two more source-level fixes in the same script: + +- **`nn/rnn_state.py` emits ill-formed TIR.** It binds the recurrent-state gather/scatter indices to + annotated locals (`seq_id: T.int32 = seq_slot_ids[vi]`); this TVM then infers a block read region + whose bound names a buffer `seq_id` that was never declared, and the **C++ verifier** rejects it — + so silencing the TVMScript parser's `check_well_formed` only defers the failure to a later pass. + Inlining each load at its single use site is the same computation and verifies clean. Qwen3.5 is + 75% GatedDeltaNet, so RNNState is on the critical path and cannot be skipped. +- **Every loader passes `mlc_param.dtype` into `numpy.astype`**, but here that is a + `tvm.ir.type.PrimType` (`T.float16`), not the `str` numpy needs. `str()` yields exactly `float16`. + +**5. `--device webgpu` needs `web/dist/wasm/*.bc`, which the wheels neither ship nor can build.** +Their header tree is pruned to inconsistency (`include/tvm/` has only `runtime/`, yet +`runtime/tensor.h` includes the absent `tvm/support/io.h`). The runtime is therefore built from an +`apache/tvm` checkout pinned to **the commit the wheel records** (`tvm._version.__commit_id__`, +here `837cb9de1`) — the linked runtime must come from the same commit as the codegen that produced +the kernels. Two further wrinkles: + +- **Build it outside a path containing spaces.** TVM's `web/Makefile` interpolates `TVM_ROOT` into + unquoted `-I` flags, so "Everything WebGPU" splits and clang fails on phantom directories. Hence + `~/.cache/mlc-web-build`. +- **`mlc_wasm_runtime.cc` is now only `#define`s** — its actual runtime moved into TVM's + `wasm_runtime.cc` — so it compiles standalone with no include path at all. + +### And the one that only shows up at runtime + +TVM's `web/emcc/wasm_runtime.cc` pulls in tvm-ffi's sources one file at a time and **misses +`custom_allocator.cc`**. `TVMFFIGetCustomAllocator` is then left undefined and becomes a wasm +*import* — and no published `@mlc-ai/web-llm` supplies it (0.2.84 is latest; there is no newer +version to upgrade to). The model compiles, ships, ingests, and then dies at instantiation: + +``` +FAIL: load: LinkError: import object field 'TVMFFIGetCustomAllocator' is not a Function +``` + +Adding that one `#include` fixes it. The check that matters is not "does it compile" but **does the +import list match a known-good build** — compare against the reference wasm: + +```sh +node tools/wasm-imports.mjs dist/Qwen3.8-2B-q4f16_1-MLC/*.wasm \ + ~/Downloads/Qwen3.5-0.8B-q4f16_1-MLC/*.wasm +``` + +Pass a reference as the second argument and it exits non-zero on any import the reference does not +need. This build now asks for 10, all of them in the reference's 11 (it does not need `proc_exit`). +Anything *extra* means the runtime is newer than the JS glue, and it will fail only at load time. + +TVM also links the final `.wasm` with `emcc`, which under emscripten 6.x no longer pulls in the C++ +stdlib for bitcode inputs — `wasm-ld` fails on `std::cerr` and the `ostringstream` vtables. Patched +to `em++`, which is emscripten's own suggestion in the error. + +## Resolved unknowns from the original handoff + +| question | answer | +| --- | --- | +| keep `model.language_model.` prefix or rewrite? | **Keep it.** `qwen35_loader.py` hardcodes `hf = "model.language_model"` and maps MLC `model.X` -> `model.language_model.X`. Use `--keep-prefix`. | +| is flattening `text_config` the step most likely to fail? | **No — it is handled upstream.** `Qwen35Config.__post_init__` already unwraps both `text_config` and nested `rope_parameters`. The strip script still flattens, which is equivalent. | +| flags to exclude tree-attn / sliding-window / paged-prefill? | **None exist**, and none can — fixed positional tuple, see above. | +| cost of no KV reuse? | **5.27 ms per history token**, floor 624 ms. Measured, and it matters. | +| does MLC support this architecture? | Yes — `qwen3_5` in the registry, `mlc_llm/model/qwen35/`, hybrid `kv_state_kind: 'hybrid'` (paged KV + RNN state). | + +### Two things the reference 0.8B gets wrong that this build does not + +- **Stop tokens.** The reference was built with `--conv-template qwen2`, so its `stop_token_ids` are + `[151643, 151645]` — **Qwen2 ids, which do not exist in this 248320-token vocab**. It survives + only on `stop_str` text matching. This build uses `--conv-template qwen3_5`, whose template + already carries the correct `[248046, 248044]` (`<|im_end|>`, `<|endoftext|>`), and `gen_config` + independently picked the same pair out of `generation_config.json`. +- **Thinking.** `qwen3_5` opens the assistant turn with ``, which is what this checkpoint + expects — the model card says every answer opens with a `` block. Confirmed in the e2e + output. Use `qwen3_5_nothink` to suppress it. + +## Notes on the checkpoint + +`empero-ai/Qwen3.8-2B-Distill`, at `model/Qwen3.8-2B-Distill/` (gitignored). **Measured** from the +safetensors header: 632 tensors, 2.27 B params, cleanly separated by prefix — +`model.language_model` (320 / 1.88 B, keep), `model.visual` (297 / 331 M, drop), `mtp` (15 / 61 M, +drop). The model card confirms the fine-tune is text-only with vision inherited from the base. + +`tools/strip-vision.py` streams byte ranges out of the source safetensors rather than loading it, so +it runs in a few MB rather than the 4.5 GB the tensors occupy — worth keeping on a 16 GB machine. +It also flattens the config and hoists `rope_theta` / `partial_rotary_factor`, and folds both EOS +ids from `generation_config.json` into `eos_token_id`. + +Architecture vs the verified 0.8B: only `hidden_size` (1024 -> 2048) and `intermediate_size` +(3584 -> 6144) differ. Layers 0-2 of each group of 4 are GatedDeltaNet linear attention +(`A_log`, `dt_bias`, `conv1d`, `in_proj_{qkv,a,b,z}`); every 4th is full attention with a gated +`q_proj` of `[4096, 2048]` = 2 x 8 heads x 256 (`attn_output_gate: true`). + +**Keep `q4f16_1`.** Dequantisation is ~free (41 GB/s with unpack-and-scale vs 41-46 without), so a +wider format only adds bytes, and bytes are what decode pays for. + +### Sampling + +The model card recommends `temperature=0.6, top_p=0.95, top_k=20`, and warns that greedy decoding is +a repetition-loop failure mode for this class. `gen_config` writes MLC's defaults (1.0 / 1.0), so +**this build overwrites two of the three** in `mlc-chat-config.json`: + +| | value | effective? | +| --- | --- | --- | +| `temperature` | 0.6 | **yes** — the engine default was moved to match | +| `top_p` | 0.95 | **yes** — nothing overrides it | +| `top_k` | 20 | **no** — compiled in, but unreachable | + +`gen_config` rewrites this file, so a rebuild resets both. Re-apply after it: + +```sh +python3 - <<'EOF' +import json, pathlib +p = pathlib.Path("dist/Qwen3.8-2B-q4f16_1-MLC/mlc-chat-config.json") +c = json.loads(p.read_text()); c["temperature"] = 0.6; c["top_p"] = 0.95 +p.write_text(json.dumps(c, indent=2) + "\n") +EOF +``` + +**`top_k` has nowhere to land.** The library genuinely supports it — `get_renorm_prob` takes +`top_k: Tensor([batch, 1], int32)` and the WGSL applies a real top-k cutoff beside top-p — but +`top_k` appears **zero** times in `vendor/web-llm.js`, `src/background/multistep.js` and +`src/background/engine-worker.js`. The burst sampler calls `fsampleWithTopP` with temperature and +top_p only. So it is the same shape of problem as the 10-buffer kernels: present in the file, dead +on the live path. Reaching it means wiring a `top_k` through *both* the multistep burst sampler and +the vendored WebLLM path — partial wiring would make burst and non-burst decoding sample +differently, which is worse than leaving it. The reference 0.8B is in exactly the same position. + +**`temperature` needed a second change.** `src/background/background.js` sets +`temperature: payload.temperature ?? settings.temperature` on every request, so +`DEFAULT_SETTINGS.temperature` in `src/lib/model-store.js` shadows whatever the model ships — +setting `mlc-chat-config.json` alone would have done nothing. That default is now `0.6` to match. +It is a single global, not per-model, so it applies to any other model loaded in the extension. +`top_p` is never injected by `buildParams`, which is why 0.95 takes effect straight from the model +config. + +## Verify + +```sh +node tools/audit-wasm.mjs dist/Qwen3.8-2B-q4f16_1-MLC # exits 1: see the buffer section above +node tools/wasm-imports.mjs dist/Qwen3.8-2B-q4f16_1-MLC/*.wasm +MODEL_DIR="$PWD/dist/Qwen3.8-2B-q4f16_1-MLC" ENGINE_COUNT=1 npm run e2e +``` + +The ingest contract (`src/lib/ingest.js`) is satisfied: `mlc-chat-config.json`, `tensor-cache.json` +with 26 `dataPath` records all present, `tokenizer.json` both present and listed in +`tokenizer_files`, and exactly one `.wasm`. `gen_config` also emits `tensor-cache-b16.json` (same 26 +shards); it is ignored by ingest and harmless. diff --git a/README.md b/README.md new file mode 100644 index 0000000000000000000000000000000000000000..5b74fe8f849cfd0fef4e9666a42c11c763ef57e5 --- /dev/null +++ b/README.md @@ -0,0 +1,496 @@ +--- +license: isc +library_name: everything-webgpu +pipeline_tag: text-generation +tags: + - webgpu + - mlc-ai + - web-llm + - browser + - on-device + - local-llm + - ghost-text + - scheduler + - typescript + - javascript +--- + +# everything-webgpu + +[![npm version](https://img.shields.io/npm/v/everything-webgpu.svg?color=cb3837)](https://www.npmjs.com/package/everything-webgpu) [![GitHub](https://img.shields.io/badge/GitHub-nyaaorick%2Feverything--webgpu-black?logo=github)](https://github.com/nyaaorick/everything-webgpu) [![WebGPU](https://img.shields.io/badge/Platform-WebGPU-orange.svg)](https://developer.mozilla.org/en-US/docs/Web/API/WebGPU_API) [![License: ISC](https://img.shields.io/badge/License-ISC-blue.svg)](https://opensource.org/licenses/ISC) + +Run local MLC/WebLLM models on WebGPU, in a browser, with a scheduler in front of them. + +**Migrating off `@mlc-ai/web-llm` is one line.** + +```diff +-import { CreateMLCEngine } from "@mlc-ai/web-llm"; +-const engine = await CreateMLCEngine(modelId, { initProgressCallback }); ++import { CreateScheduledEngine } from "everything-webgpu"; ++const engine = await CreateScheduledEngine(modelId, { initProgressCallback }); +``` + +Everything after that line is unchanged. `engine.chat.completions.create()` takes and returns the +same shapes, streams the same chunk objects, and reports the same finish reasons. That call is the +compatibility layer and **it never changes** — see [Stability](#stability). + +What the swap buys: + +| | | +| --- | --- | +| **A scheduler** | `session`, `priority`, `task`, `preemptible` on any request. One engine per task, pool grows on demand, stale requests supersede themselves. | +| **Multi-step decoding** | K forward steps per GPU sync. On the 0.8B baseline: 9.6 → 25.9 tok/s. | +| **Two build-time patches** | Firefox's 9-storage-buffer cap, and one compute pass per flush instead of per kernel. | +| **Three model sources** | A prebuilt id, a URL you host, or **a folder off disk with no network at any point**. | +| **Typed errors** | Eight codes, each one because a caller does something different about it. | + +--- + +## Status + +**On npm** — `npm i everything-webgpu`. Pre-1.0, ESM only (`"type": "module"`). + +```sh +npm i everything-webgpu +``` + +`vendor/web-llm.js` — the patched WebLLM bundle the engine loads — is a **build product and is not +checked in**, so every install route has to produce it. The published tarball ships it prebuilt; +`prepare` regenerates it for a git dependency or a clone. If it is ever missing you get +`PACKAGE_INCOMPLETE` naming the fix, rather than a module-resolution error pointing inside this +package. + +Requires WebGPU. Verified end to end on Firefox 154 / macOS / M4; on Firefox you may need to set +`dom.webgpu.enabled` and the related flags first — [AI.md](AI.md) lists them and says what each one +does. Chrome is expected to be faster (it does not hit the storage-buffer cap that disables KV +reuse) but that is a prediction, not yet a measurement. + +## Quick start — four lines + +```js +import { CreateScheduledEngine } from "everything-webgpu"; // 1. import +const engine = await CreateScheduledEngine("Llama-3.2-1B-Instruct-q4f16_1-MLC"); // 2. load a model +const reply = await engine.ask("Name three primary colours."); // 3. ask +console.log(reply); // 4. the answer +``` + +`reply` is a plain string. Line 2 downloads ~0.8 GB the first time and prints throttled progress to +the console; pass `initProgressCallback` to render it yourself, or `initProgressCallback: null` to +silence it. On Vite, add one plugin — see [API.md § Bundlers](API.md#bundlers). Every other call +form — `conversation()`, `load()` from a URL or off disk, streaming, embeddings, the OpenAI shape — +is catalogued in **[API.md](API.md)**. + +The WebLLM-compatible shape, unchanged: + +```js +const reply = await engine.chat.completions.create({ + messages: [{ role: "user", content: "hi" }], + session: "ghost-text", // added by this engine + priority: "interactive", // added by this engine +}); +console.log(reply.choices[0].message.content); +``` + +Streaming is WebLLM's, passed through untouched: + +```js +const stream = await engine.chat.completions.create({ messages, stream: true }); +for await (const chunk of stream) process.stdout.write(chunk.choices[0]?.delta?.content ?? ""); +``` + +The first `load()` downloads the weights; after that it is cache-only and needs no network. + +### Runnable + +[examples/](examples/) has three, each a standalone project that consumes this package through its +`exports` map rather than reaching into the source tree: + +```sh +cd examples/bare && npm install && npm run dev +``` + +| | | +| --- | --- | +| [bare/](examples/bare/) | Vanilla Vite, one page. ~100 lines. | +| [react/](examples/react/) | A chat UI over `conversation()`. | +| [webext/](examples/webext/) | MV2 Firefox — `browser.storage.local` as the store, and the CSP a wasm runtime needs. | + +## What it costs your bundle + +**53 kB** (~19 kB gzip) in the entry chunk, before a model is loaded. WebLLM's 6 MB bundle is a +**lazy chunk** — fetched on the first `load()` or `listAvailableModels()`, and never by a visitor +who does neither, so a page can decide whether this machine has a GPU before paying for it. The +IndexedDB adapter is another 0.8 kB lazy chunk, absent entirely if you pass your own store. + +On Vite, add the plugin this package ships: + +```js +import { everythingWebGPU } from "everything-webgpu/vite"; +export default defineConfig({ plugins: [everythingWebGPU()] }); +``` + +It exists for one reason: Vite's dependency pre-bundler otherwise rewrites the +decode worker's URL to a path that 404s — in `vite dev` only, and only on a real +install, which is why no example in this repo can catch it and +`npm run verify-consumer` exists instead. Skip it and `load()` fails with +`PACKAGE_INCOMPLETE` naming the fix, rather than hanging. + +## Three ready-made shapes + +Most apps want one of three things. Each is one call, with the scheduling already right: + +```js +// 1. One question, nothing kept. +const answer = await engine.ask("Summarise this in one line:\n" + doc); + +// 2. A conversation that remembers. +const chat = engine.conversation({ system: "You are terse." }); +await chat.say("what's the capital of France?"); +await chat.say("and its population?"); // remembers + +// 3. Ghost text, debounced and superseding. +const ghost = engine.ghostText({ prompt: (before) => `Continue:\n${before}` }); +editor.on("input", async () => { + const hint = await ghost.suggest(editor.textBefore()); + if (hint !== null) render(hint); // null = a newer keystroke won +}); +editor.on("blur", () => ghost.cancel()); +``` + +`complete()` can express all three. These exist because the **scheduling** is the part that's easy +to get wrong and invisible when you do — every row of [Getting these wrong](#getting-these-wrong) is +a scheduling mistake, not a generation one. So each verb is a policy: + +| | policy | +| --- | --- | +| `ask()` | its own task, no session — two in flight never supersede each other | +| `conversation()` | one stable task for every turn, turns serialised, history bounded | +| `ghostText()` | debounce + one session key + `interactive` + **resolves `null` when stale** | + +That last one is the difference between a verb and a wrapper. The engine already drops superseded +work; what a caller still had to remember was not to *paint* the answer that came back anyway. +Returning `null` removes the choice. + +**They author no prompts.** `ask()` and `conversation()` carry your text through; `ghostText()` +*requires* a `prompt` function and has no default. Prompts are model-specific — switching this +project's own build from a 0.8B to a 2B changed the conversation template and made every reply open +with a `` block. A prompt that lives in your code survives that. + +`conversation()` bounds history at 12 exchanges by default, because there is no cross-turn KV reuse +here: every turn re-prefills the whole history at ~5.27 ms/token, so an unbounded conversation gets +quadratically slower and a turn near the context limit waits ~22 s for its first token. Pass +`keep: Infinity` to opt out, having read that sentence. + +## Where models come from + +**One call covers all three.** `load()` works out what you handed it, registers whatever needs +registering, and brings the model up: + +```js +// 1. Prebuilt — one of WebLLM's 163 HuggingFace models, or anything you registered earlier. +await engine.load("Llama-3.2-1B-Instruct-q4f16_1-MLC"); + +// 2. Remote — any base URL you host. `modelLib` is required; see below. +await engine.load("https://cdn.example/models/my-model/", { + modelLib: "https://cdn.example/models/my-model/my-model-webgpu.wasm", +}); +await engine.load({ model: "/models/my-model/", modelLib: "/models/my-model/lib.wasm" }); + +// 3. Injected — a folder off disk. No network at any point. +await engine.load(dropEvent.dataTransfer); // a drop +await engine.load(input.files); // +await engine.load({ files }); // already-unpacked { path, file }[] +``` + +The model id is derived from the URL's last segment; pass `{ id }` to override it. Add +`{ defer: true }` to register a source **without** building a pool — it returns the registry record, +for a drop-now-load-later flow. `registerModel()` and `ingestModelFolder()` remain as the low-level +primitives; `load()` composes them rather than replacing them. + +`listAvailableModels()` enumerates all three routes. Construct the engine with `{ prebuilt: false }` +for a build that must never fetch a model over the network. A misspelled id is told what it might +have meant. + +### Warming the cache early + +```js +await engine.prefetch(modelId, { onProgress: (p) => bar(p.done / p.total) }); +``` + +Downloads the model **without building an engine and without WebGPU** — so an app can warm the cache +while the user is still reading the welcome screen, before it has decided whether this machine can +run the model at all. WebLLM cannot express this: `reload()` instantiates the wasm and needs a GPU +before it fetches a single shard. A later `load()` is then a cache read. Interrupted downloads +resume, and a second call is free. + +## Embeddings + +```js +await engine.load("snowflake-arctic-embed-s-q0f32-MLC", { keepResident: true }); + +const [q] = await engine.embed("how do I cancel a job?"); +const docs = await engine.embed(paragraphs); // one vector per input, in order +``` + +Embedding models are separate models — WebLLM ships four `snowflake-arctic-embed-*` from 239 MB — so +this usually names `modelId` and holds it resident alongside a chat model. Embeddings go through the +**same scheduler** as completions: same priority bands, same `session` supersession, same +one-task-one-engine rule. `embedRaw()` returns WebLLM's OpenAI-shaped envelope if you are porting +code that expects `data[].embedding`. + +One difference worth knowing: **a running embedding cannot be interrupted.** Cancellation works by +making a decode loop break out, and one forward pass has no loop — so a `cancel()` landing after the +job starts marks it cancelled without stopping it. Queued embeddings supersede normally. + +**`modelLib` is required for a remote source and is not guessed.** It is genuinely underivable: +across all 163 prebuilt models, zero have a lib name derivable from the model id, and zero host the +lib on the same origin as the weights. A guess would be wrong every time and would surface as a +confusing 404 deep inside the loader, so the engine asks instead. + +## The three shapes of work + +What differs between these is *not* the call or the transport. It is **who owns an engine, and what +may interrupt what.** + +| | call | priority | key field | why | +| --- | --- | --- | --- | --- | +| **Completion** (ghost text) | `complete` streaming | `interactive` | `session` | Each keystroke supersedes the last request; may preempt opted-in work. | +| **Translation** (a page) | `batch` | `normal` | one shared `task` | One request instead of N, so the engine schedules it as a unit and it never hogs the pool. | +| **Reformat** (markdown) | `complete` | `background` | `preemptible: true` | Nobody is watching; let interactive work cut in. | + +### Completion — latency is the whole product + +```js +// On every keystroke. The previous request is cancelled, not queued behind. +const { text } = await engine.complete( + { + messages: [{ role: "user", content: prefix }], + session: "ghost-text", // supersession key — the important field + priority: "interactive", // may preempt jobs that opted in + max_tokens: 24, // ghost text is short; do not pay for more + }, + (delta) => render(delta), +); +``` + +**`session` is what makes this work, not `cancel`.** Reusing one session key means the engine drops +the stale request itself. A caller that mints a fresh id per keystroke and calls `cancel` races its +own typing. + +### Translation — throughput, one task + +```js +// One batch, not a loop of `complete` calls. +const results = await engine.batch({ + task: "translate-page", // optional; a batch is one task either way + requests: sentences.map((s) => ({ + messages: [{ role: "user", content: `Translate to French, output only the translation:\n${s}` }], + })), +}); +results.forEach((r) => apply(r.index, r.text)); +``` + +Every item of one batch shares a task, and a task holds one engine, so a 200-sentence page occupies +exactly one engine and can never freeze ghost-text behind it. Items carry `engineIndex`, `startedAt` +and `finishedAt`, so you can check what actually ran where. + +`batch` stays the right call rather than a loop: it is one scheduling unit the engine can reason +about, and when batched decode lands the same call gets faster with no change on your side. + +### Reformat — cheap to interrupt + +```js +await engine.complete({ + messages: [{ role: "user", content: `Reformat as clean Markdown, no commentary:\n\n${doc}` }], + priority: "background", + preemptible: true, // the direction matters — see below + max_tokens: 2048, +}); +``` + +**Set `preemptible` on the work that can afford to lose, not on the work you care about.** Only an +`interactive` request preempts, and only a job that opted in can be preempted. A preempted job +resolves with `preempted: true` and whatever text it had — never requeued, so it can never starve, +but you must be able to use or discard a partial result. + +## Getting these wrong + +| symptom | cause | +| --- | --- | +| Ghost text lags behind typing | Fresh `id` per keystroke with no `session`, so every stale request still runs. | +| Page translation is slower than expected | Expected: one task is one engine, and a second engine measured 1.06x anyway. Throughput here comes from batched decode, not from more engines. | +| Reformatting blocks completions | `preemptible` left off the background job, so `interactive` has nothing to take. | +| Pool stays at one engine | Expected: it grows only when a *second task* waits. Check `engine.state.pool.growthBlocked` if two are waiting and it still has not. | + +## Prompts stay with you + +The engine does not author prompts. They are model-specific — switching one build from +`Qwen3.5-0.8B` to `Qwen3.8-2B-Distill` changed the conversation template and made every reply open +with a `` block. A prompt that lives in the caller survives that; a `translate` op baked into +the engine would have to be rewritten and re-shipped to every caller. + +## When it does not work + +Every failure carries a code. `message` stays human-readable and stays the thing you print; `detail` +carries structured context so you never parse the sentence. + +```js +import { isEngineError, ERROR } from "everything-webgpu"; + +try { + await engine.load(id); +} catch (err) { + if (isEngineError(err, ERROR.CACHE_INCOMPLETE)) return reRegisterFolder(); + throw err; +} +``` + +| code | what to do about it | +| --- | --- | +| `NO_WEBGPU` | Tell the user to check flags/hardware; retrying is futile. | +| `NO_MODEL` | Nothing registered at all — send them to your setup flow. | +| `UNKNOWN_MODEL` | That id is not resolvable; `listAvailableModels()` says what is. | +| `CACHE_INCOMPLETE` | A locally-registered model was evicted; re-register the folder. | +| `INVALID_MODEL_FOLDER` | Not a compiled MLC model; `detail` says what is missing. | +| `BAD_REQUEST` | The caller's arguments are wrong — a bug in the caller. | +| `ABORTED` | The caller cancelled it. Not a failure; do not report it as one. | +| `GENERATION_FAILED` | The model failed mid-generation. | +| `PACKAGE_INCOMPLETE` | Your **build** is wrong, not your code — a missing `vendor/` bundle or a decode worker the bundler did not emit. `message` names the fix; `detail.cause` says which. | + +### Ask before you download, not after + +`environment()` is the preflight: one call that says what this machine will do, and why. + +```js +const report = await engine.environment(); +if (!report.ok) console.error(report.lines[0].fix ?? report.lines[0].cause); +``` + +Every line carries the same five fields, so you can render the whole report without special-casing +any of it: + +| field | | +| --- | --- | +| `severity` | `blocked` · `degraded` · `tune` · `info` · `ok` — lines come sorted worst-first | +| `affects` | what you lose, in your terms | +| `cause` | the measured fact behind the verdict | +| `fix` | the exact call to make, or `null` when nothing can be done | +| `operable` | whether this is reachable from JS at all | + +`fix` is `null` for hardware, build-time flags, and browser settings JS cannot reach. Reporting a +consequence with no remedy is still the point: *"your second turn is slow because this device caps +storage buffers at 9"* is the difference between a bug report and an informed decision. + +```js +await engine.environment({ scope: "local" }); // never touches the model layer — cheap, poll freely +await engine.environment({ scope: "device" }); // hardware only +await engine.environment.measure(); // one calibration generation → measured tok/s +``` + +**`environment()` only reports.** Writes go through `configure()`, and passing a setting to +`environment()` is an error that names the call you wanted — one function doing both, told apart by +argument shape, is how you get a call that silently does nothing. + +```js +await engine.configure({ decodeSteps: 8 }); // hot, no reload +await engine.configure({ engineCount: 3 }); // persisted; applies to pools built after it +``` + +`canRun(modelId)` stays separate and answers the per-**model** question — will this one fit and run +here — where `environment()` is per-**device**. + +## Freeing memory + +A resident model is a full copy of its weights in VRAM, and nothing reports free VRAM to a page — so +residency is explicit rather than guessed at. + +```js +await engine.unload(); // the current model's VRAM; cached bytes stay, so reloading is free +await engine.unload(id); // that model's VRAM +await engine.unload(id, "cache"); // and delete its cached bytes, keeping the registry entry +await engine.unloadAll(); // every resident model +await engine.remove(id); // bytes and entry; for an injected model, means re-supplying the folder +``` + +A bare `unload()` frees **only the current model** — `unloadAll()` is the explicit form, because +freeing everything is the more destructive reading and should have to be asked for by name. +`remove()` keeps its own verb for the same reason: it is the one that cannot be undone. +`evict(id)` remains as the low-level primitive `unload(id, "cache")` is built on. + +`load()` unloads whatever else is up before bringing a model in, which is the safe default on a +16 GB machine. Pass `{ keepResident: true }` to hold two at once, having checked the budget with +`canRun()` first. `cacheState(modelId)` says what is already on disk. + +## Choosing a store + +`CreateScheduledEngine` defaults to IndexedDB, because a registry that dies with the page would +strand the weights — the bytes stay in Cache Storage but nothing remembers they are there. Pass one +explicitly in a worker or a test: + +```js +import { ScheduledEngine, ModelStore } from "everything-webgpu"; +import { memoryStorage } from "everything-webgpu/adapters/memory"; + +const engine = new ScheduledEngine({ store: new ModelStore(memoryStorage()) }); +``` + +Adapters ship for `idb` (pages, plus `ensurePersistent()`), `memory` (tests), and `webext` +(`browser.storage.local` and the message/port router). + +## Measured + +M4 MacBook Air (16 GB), Firefox 154 release, macOS. `Qwen3.8-2B-q4f16_1` (1.06 GB) shipping, +`Qwen3.5-0.8B-q4f16_1-MLC` (443 MB) as the baseline most of the analysis was done on. + +| | 2B | 0.8B | +| --- | --- | --- | +| Model load, cache only, zero network | 51 s | 48 s | +| Prefill | 48 tok/s short, 100–200 at length | 95–98 tok/s | +| **Decode** | **16.6–18.1 tok/s** | **25.9 tok/s** | +| Decode, stock single-step WebLLM | — | 9.6 tok/s | +| Second engine, 4-prompt batch | 1.06x | — | + +Reproduce with `npm run e2e`. The full analysis — where every millisecond goes, why a second engine +buys isolation rather than throughput, and what is still open — is in [AI.md](AI.md). + +## Stability + +`chat.completions.create()` is the compatibility layer. **It does not change**, and any change to +its request or response shape is a major version. That is the whole point of the one-line migration: +if it drifted, the line would not be one line. + +Everything else is pre-1.0 and may move. The ergonomic verbs (`load`, `unload`, `environment`) are +being consolidated — see [ROADMAP.md](ROADMAP.md). + +## Docs + +| | | +| --- | --- | +| [API.md](API.md) | Every way to call it — one page, asserted against the code. | +| [examples/](examples/) | Three runnable projects. Read one before reading the source. | +| [AI.md](AI.md) | What is true and measured. The reference. | +| [ROADMAP.md](ROADMAP.md) | The only list of open work. | +| [ARCHIVE.md](ARCHIVE.md) | What was done and *why* — decisions with their reasoning. | +| [WEBLLM-SURFACE.md](WEBLLM-SURFACE.md) | What WebLLM already does, where our line is, and the dependency-bump runbook. | +| [MLC-COMPILE.md](MLC-COMPILE.md) | How the model was compiled, and every toolchain breakage on the way. | + +## Development + +```sh +npm run build # bundle WebLLM + apply the two patches +npm test # contract, patches, README, examples, scheduler, multistep — GPU-free +npm run e2e # real Firefox + real GPU + real model +npm run verify-consumer # pack + install for real; proves the Vite plugin is needed and works +npm run verify-patches # check patch anchors without rebuilding +``` + +`@mlc-ai/web-llm` is pinned exactly, because the build rewrites its bundle at verified anchors. +Bumping it is a documented, guarded procedure — see WEBLLM-SURFACE.md, "Upgrading". + +## Licence + +This package is **ISC** ([LICENSE](LICENSE)). It ships a pre-built copy of WebLLM +(`vendor/web-llm.js`, **Apache-2.0**) and its dependency `loglevel` (**MIT**); their full licence +texts travel in [THIRD-PARTY-NOTICES.md](THIRD-PARTY-NOTICES.md), and +[test/license.test.mjs](test/license.test.mjs) fails the build if a dependency is ever bundled +without one. diff --git a/ROADMAP.md b/ROADMAP.md new file mode 100644 index 0000000000000000000000000000000000000000..cd891c43ad7ee910572c809a231bce11eac24297 --- /dev/null +++ b/ROADMAP.md @@ -0,0 +1,181 @@ +# Roadmap — open work only + +**This is the only list of open work.** Anything finished moves to [ARCHIVE.md](ARCHIVE.md) with the +reasoning that produced it. [AI.md](AI.md) is the reference for what is true and measured; +[API.md](API.md) is the full call surface, asserted against the code; +[WEBLLM-SURFACE.md](WEBLLM-SURFACE.md) is what WebLLM already does and must be read before adding a +capability. + +Sections are named, not numbered. An earlier split across three files used "Track 1"/"Track 2" to +mean different things in each, which is how a reader ends up implementing the wrong item. + +**Sequencing.** The whole project exists to make WebLLM *more compatible, easier to use, foolproof +to build on*. Every item below is ranked against that. `everything-webgpu@0.1.0` is on npm as of +this writing — the install path, the four-line demo, the ergonomic verbs, `API.md` and `examples/` +all shipped and are recorded in [ARCHIVE.md](ARCHIVE.md). Section 1 is what remains before the +extraction is fully proven: the extension itself rebuilding on the package. Section 4 (performance) +is frozen as a block, except for "Measure Ollama", until a real user calls it slow: it is the +highest-risk, highest-effort work and it re-touches the tvmjs internals the WebLLM-upgrade section +(also in [ARCHIVE.md](ARCHIVE.md)) just spent its time hardening. + +--- + +## 1. Ship it — the last mile + +`everything-webgpu@0.1.0` is on npm. What shipped — the install path, the four-line demo, the +ergonomic verbs, [API.md](API.md), [examples/](examples/), the bundle-size story, licence +compliance — is in [ARCHIVE.md](ARCHIVE.md), "Shipping 0.1.0". What remains is proving the boundary +from the other side. + +- [ ] **`demo` consumes the package**, and the popup + manager leave `main`. The acceptance test for + the whole extraction: if the extension rebuilds on the published package, the boundary is + right. `npm run e2e` passing is the same claim at the source-tree level. +- [~] **Gate A — a bare Vite page loads a prebuilt model and generates.** + [examples/bare/](examples/bare/) exists and the GPU-free half is proven: `vite build` resolves + through the `exports` map and splits into an entry chunk, the decode worker as its own chunk, + and the 6 MB WebLLM bundle as a lazy one. What remains is **one run in a WebGPU browser** to + confirm it generates. `npm run verify-consumer` already exercises a real install end to end, + short of the actual generation. +- [ ] **Gate A′ — the same page ingests a local folder.** `cache.put()` against the synthetic + `local-model.invalid` key is proven only on an extension origin. Gates the *offline* route, + not the library. +- [ ] **Gate B — Chrome.** Measure tok/s. Expected to beat Firefox because KV reuse is not disabled + there. Entirely a prediction today; `probeDevice()` is the instrument that makes it reportable. +- [ ] **Every `ERROR` code carries an actionable `fix`.** The preflight (`environment()`) and the + typed errors (9 codes now, `PACKAGE_INCOMPLETE` included) both exist; what is not yet done is + the systematic pass confirming each code's `message`/`detail` names a cause *and* a fix, and + that the README/API.md point at `environment()` as the "will this run here?" call it is meant + to be. + +## 2. Verb consolidation — done + +`load(src, opts)`, `unload(id, level)` / `remove()` / `unloadAll()`, and `environment()` + +`environment.measure()` all shipped. The reasoning — the two dropped `load()` dispatch rules, why a +bare `unload()` frees only the current model, why read and write are split — is in +[ARCHIVE.md](ARCHIVE.md), "Verb consolidation". `chat.completions.create()` is unchanged and stays +that way. + +## 3. Engine capability + +`prefetch()`, `embed()` / `embedRaw()`, and the `ask()` / `conversation()` / `ghostText()` recipes +shipped — see [ARCHIVE.md](ARCHIVE.md), "Engine capability". One item is still open: + +- [ ] **LRU eviction** when quota is short. `cacheState()` and `evict()` exist; nothing yet decides + *what* to drop. Table stakes for "it just works" the moment a second model is cached — but the + eviction policy is a decision, not an implementation detail. + +## 4. Performance — frozen, except the measurement + +Model- and kernel-level work, independent of the library structure — and therefore independent of +everything this project is *for*. **Frozen as a block until a real user calls it slow.** The library +ships now, so "nobody can install it" is no longer the argument; the argument is that this is the +highest-risk, highest-effort work on the page, it re-touches the tvmjs internals the WebLLM-upgrade +section (in [ARCHIVE.md](ARCHIVE.md)) just spent its time hardening, and a speedup nobody has asked +for is risk without a return. The one exception is "Measure Ollama": it is cheap, and it decides +whether any of the rest is worth its risk. + +- [ ] **Retune the dlight GEMV schedule** — more work per thread before the reduction. Measured + **1.83x** in isolation, which would put decode near the ~41 GB/s dequant ceiling. Confirmed + untaken on the 2B build. Needs a recompile; toolchain is stood up. +- [ ] **Batched decode.** The model lib exports `batch_decode` / `batch_prefill` / `batch_verify` + and a paged KV cache; WebLLM hardcodes `defaultMaxNumSequence = 1` and `numSamples = 1` + (bundle lines 15250, 15276). Lifting that reads the weights once per step for N sequences — + projected **~4x on the `batch` API**, which is the page-translation shape and the one workload + the engine currently makes no faster than a loop of `chat` calls. No recompilation needed. + **Not** a fix for the 1.06x second-engine measurement: a second engine buys task isolation, + never aggregate throughput, and the two are complementary rather than alternatives. +- [ ] **Interrupt granularity in `sampleBurst`.** `interruptGenerate()` only sets a flag the + caller's loop reads between `decode()` calls, and a burst runs all K forward steps without + checking it — so a preempted job finishes its whole burst first (~583 ms at K=15). It binds in + three cases: the seconds-long window while `#grow()` builds a second engine, a machine where + `#growthBlocked` is set, and any third concurrent task once the pool is at cap. Fix in our + `multistep.js`, not WebLLM; `discardLookahead` already handles the resulting state. +- [ ] **Auto-tune `decodeSteps` from the burst probe.** K=15 is a constant derived from one machine + and one model — Firefox on Apple Silicon, `Qwen3.5-0.8B` at ~7.3 ms/step. `multistep.js`'s own + header says the best K falls as the model grows (a 25 ms/step model wants K=4), so the shipped + default is wrong for every model but the one it was measured on, and wrong silently: a bad K + costs throughput and reports nothing. + + **The infrastructure is already there.** `onBurst` delivers `{ steps, tokens, ms, encodeMs, + syncMs, dispatches, forwardDispatches, flushes }` per burst + ([multistep.js:246](src/engine/multistep.js#L246)), `setSteps()` is hot + ([multistep.js:259](src/engine/multistep.js#L259)), and `WORKER_CONFIGURE` already calls + `resetStats()` on every retune so "a sweep's points never bleed into each other" + ([engine-worker.js:137](src/engine/engine-worker.js#L137)). It was built for this. + + **The obvious derivation does not work.** `K = floor(tickMs / perStepMs)` needs `perStepMs`, and + the probe cannot separate it: `syncMs` is GPU execution *plus* the wait for the poll tick, which + is the quantity being solved for. Only `encodeMs` is clean — the K-step loop contains no `await`, + so it is content-process CPU alone — and it is the smaller half. + + **What works instead** is a hill-climb on the one quantity directly observed, `tokens / ms`. + Throughput at K is `K / (ceil(K·p/tick)·tick)`; its peaks sit at `K = floor(n·tick/p)` and are + asymptotically equal in n, so the first peak is the target and climbing K until throughput drops + finds it without ever naming `p` or `tick`. That is also what makes it survive Chrome, where the + 100 ms tick premise does not hold. + + **Three subtleties that will bite.** (1) `onBurst.steps` is the *clamped* burst size — + `burstSize()` cuts it by `max_tokens` and the context window + ([multistep.js:297](src/engine/multistep.js#L297)) — so the tuner must learn only from bursts + where `steps === config.steps`, or it learns from end-of-generation stubs. (2) The first burst on + a pipeline carries shader compilation; skip it. (3) `clampSteps("auto")` is **1** today + (`Number("auto")` → NaN → `|| 1`), so an `"auto"` sentinel silently *disables* multi-step unless + `clampSteps` and [API.md](API.md)'s documented `1`–`32` range change together — + `api-doc.test.mjs` asserts the two agree. + + **Sequencing.** Composes with per-priority K below rather than replacing it: this finds the + throughput ceiling K_max for the device+model, per-priority then spends it (`interactive` takes + min(K_max, 4), `background` takes K_max). Do this one first — per-priority is written against the + 15 this replaces. Unlike the rest of section 4 it touches **no new tvmjs internals**: it reads + telemetry that already flows and calls a setter that already exists, which is the argument for + lifting the freeze on this item alone. It cannot be *validated* here, though — the sawtooth only + exists on real hardware, and Gate B (Chrome) is what proves the controller is finding a real peak + rather than re-deriving 15. + +- [ ] **Per-priority `decodeSteps`.** K=15 maximises throughput but emits 15 tokens every ~583 ms, + which reads as a stall. `interactive` should use K=2–4, `background`/`batch` K=32. Do this + *after* the interrupt fix — raising K for background work lengthens exactly the bursts that + preemption has to wait out. +- [ ] **Restore cross-turn KV reuse.** Every turn re-prefills the whole history at **5.27 ms/token**, + so a turn near the 4096 limit pays ~22 s before its first token. Pack + `batch_prefill_paged_kv_kernel`'s six i32 metadata buffers into one with offsets (10 bindings + → 5); the offsets already exist in its uniform block. + +## 5. Test and infrastructure + +Both are "fix when it next bites", not scheduled work. The completed e2e run that used to sit here +is in [ARCHIVE.md](ARCHIVE.md), including the `storageBuffersPerStage=9` anomaly worth re-checking +on the next run. + +- [ ] Fix `PROFILE_PATH` in [test/e2e/run.mjs](test/e2e/run.mjs): it passes `--profile-path`, but + web-ext 8 calls it `--firefox-profile` and exits with `Unknown arguments`. +- [~] **Isolate the bench's pass-sweep onto its own device** — done, and it **did not do what the + item assumed**. 2048 compute passes in one encoder loses the WebGPU device, and a lost device + does not throw: later calls silently no-op. + + Two corrections from measuring it. First, the stated worry (later probes poisoned) was empty + — the sweep was already the last phase, so nothing ran after it. The real damage was to the + sweep's **own** numbers: the run that prompted this reported `perPass=-3.9us`, 512 passes + measured faster than one, which is the device dying mid-sweep and the remaining submits + becoming free. That half is fixed — `n512` now reads a plausible ~10us and `n2048` reports + `discarded (device lost during the 2048-pass encode)` instead of inventing a number. + + Second, **isolation does not contain the loss.** With the sweep on its own device, + `deviceLostDuringBench` still reports the main device lost during this phase, on both runs. A + runaway command buffer on Metal appears to reset the whole adapter rather than one device on + it. So the honest state is: sweep numbers are now trustworthy, containment is not achieved, + and the bench is safe only because this is the last phase. + + Still open, if containment is wanted: run the sweep in a dedicated worker or a separate page + so the reset cannot reach the measuring context at all. Also worth noting the attribution is + imprecise — `lostDuring` records the phase at which `device.lost` *resolves*, and the + 16384-dispatch probe immediately above is another plausible culprit. + +## 6. Deferred, with the condition that would reopen them + +| Item | Reopens when | +| --- | --- | +| **Vision / image understanding** | An in-house compact VLM exists. The prebuilt option is `Phi-3.5-vision` at 3.95 GB, projecting ~4–5 tok/s and not co-resident with a text model in 16 GB. `modelType: "vlm"` already flows through registration, so the compile lands on an engine that accepts it. Open questions to answer against the real model: whether multi-step decoding survives an image prefill, and whether the Firefox 9-storage-buffer workaround holds for the vision tower's kernels. | +| **MV3 migration** | `demo` is forced off MV2. `CreateExtensionServiceWorkerMLCEngine` solves the messaging and in-work keep-alive but **not** the actual problem — an idle service worker is killed and takes the resident multi-GB model with it. Use the native helper for messaging; hold the engine in an offscreen document, persistent page, or dedicated tab. Also verify WebGPU is exposed in a Chrome extension SW at all. | +| **SRI / `verifyIntegrity`** | Self-hosted models are hardened, or corruption is actually observed. `ModelStore.verify()` checks key *presence*; `verifyIntegrity` checks *content* — different problems. Use the native one; do not hand-roll a hash check. | +| **Multi-model via `reload([...])`** | Memory pressure shows up before scheduling contention does. Rejected for now: `reload()` is all-or-nothing, so adding a third model to `{A, B}` reloads A and B too (~51 s each). Additive residency is a hard requirement and only `#pools` provides it. | diff --git a/THIRD-PARTY-NOTICES.md b/THIRD-PARTY-NOTICES.md new file mode 100644 index 0000000000000000000000000000000000000000..8ea364e4d8ee3a3bc66b170a941124e7222a5e55 --- /dev/null +++ b/THIRD-PARTY-NOTICES.md @@ -0,0 +1,267 @@ +# Third-party notices + +`everything-webgpu` ships a pre-built, lightly-patched copy of WebLLM at +`vendor/web-llm.js` (see [WEBLLM-SURFACE.md](WEBLLM-SURFACE.md) and +[build/patches.mjs](build/patches.mjs) for what the patches are and why). That +file is a bundle: it contains the components listed below, and their licenses +travel with it. + +The patches change two code paths for Firefox's WebGPU limits and compute-pass +batching; they add no third-party code. Everything bundled is enumerated here. + +--- + +## @mlc-ai/web-llm 0.2.84 + +- License: **Apache-2.0** +- Project: https://github.com/mlc-ai/web-llm + +WebLLM's own LICENSE file (reproduced verbatim below) also carries a summary of +the components WebLLM itself vendors. + +``` + Apache License + Version 2.0, January 2004 + http://www.apache.org/licenses/ + + TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION + + 1. Definitions. + + "License" shall mean the terms and conditions for use, reproduction, + and distribution as defined by Sections 1 through 9 of this document. + + "Licensor" shall mean the copyright owner or entity authorized by + the copyright owner that is granting the License. + + "Legal Entity" shall mean the union of the acting entity and all + other entities that control, are controlled by, or are under common + control with that entity. For the purposes of this definition, + "control" means (i) the power, direct or indirect, to cause the + direction or management of such entity, whether by contract or + otherwise, or (ii) ownership of fifty percent (50%) or more of the + outstanding shares, or (iii) beneficial ownership of such entity. + + "You" (or "Your") shall mean an individual or Legal Entity + exercising permissions granted by this License. + + "Source" form shall mean the preferred form for making modifications, + including but not limited to software source code, documentation + source, and configuration files. + + "Object" form shall mean any form resulting from mechanical + transformation or translation of a Source form, including but + not limited to compiled object code, generated documentation, + and conversions to other media types. + + "Work" shall mean the work of authorship, whether in Source or + Object form, made available under the License, as indicated by a + copyright notice that is included in or attached to the work + (an example is provided in the Appendix below). + + "Derivative Works" shall mean any work, whether in Source or Object + form, that is based on (or derived from) the Work and for which the + editorial revisions, annotations, elaborations, or other modifications + represent, as a whole, an original work of authorship. For the purposes + of this License, Derivative Works shall not include works that remain + separable from, or merely link (or bind by name) to the interfaces of, + the Work and Derivative Works thereof. + + "Contribution" shall mean any work of authorship, including + the original version of the Work and any modifications or additions + to that Work or Derivative Works thereof, that is intentionally + submitted to Licensor for inclusion in the Work by the copyright owner + or by an individual or Legal Entity authorized to submit on behalf of + the copyright owner. For the purposes of this definition, "submitted" + means any form of electronic, verbal, or written communication sent + to the Licensor or its representatives, including but not limited to + communication on electronic mailing lists, source code control systems, + and issue tracking systems that are managed by, or on behalf of, the + Licensor for the purpose of discussing and improving the Work, but + excluding communication that is conspicuously marked or otherwise + designated in writing by the copyright owner as "Not a Contribution." + + "Contributor" shall mean Licensor and any individual or Legal Entity + on behalf of whom a Contribution has been received by Licensor and + subsequently incorporated within the Work. + + 2. Grant of Copyright License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + copyright license to reproduce, prepare Derivative Works of, + publicly display, publicly perform, sublicense, and distribute the + Work and such Derivative Works in Source or Object form. + + 3. Grant of Patent License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + (except as stated in this section) patent license to make, have made, + use, offer to sell, sell, import, and otherwise transfer the Work, + where such license applies only to those patent claims licensable + by such Contributor that are necessarily infringed by their + Contribution(s) alone or by combination of their Contribution(s) + with the Work to which such Contribution(s) was submitted. If You + institute patent litigation against any entity (including a + cross-claim or counterclaim in a lawsuit) alleging that the Work + or a Contribution incorporated within the Work constitutes direct + or contributory patent infringement, then any patent licenses + granted to You under this License for that Work shall terminate + as of the date such litigation is filed. + + 4. Redistribution. You may reproduce and distribute copies of the + Work or Derivative Works thereof in any medium, with or without + modifications, and in Source or Object form, provided that You + meet the following conditions: + + (a) You must give any other recipients of the Work or + Derivative Works a copy of this License; and + + (b) You must cause any modified files to carry prominent notices + stating that You changed the files; and + + (c) You must retain, in the Source form of any Derivative Works + that You distribute, all copyright, patent, trademark, and + attribution notices from the Source form of the Work, + excluding those notices that do not pertain to any part of + the Derivative Works; and + + (d) If the Work includes a "NOTICE" text file as part of its + distribution, then any Derivative Works that You distribute must + include a readable copy of the attribution notices contained + within such NOTICE file, excluding those notices that do not + pertain to any part of the Derivative Works, in at least one + of the following places: within a NOTICE text file distributed + as part of the Derivative Works; within the Source form or + documentation, if provided along with the Derivative Works; or, + within a display generated by the Derivative Works, if and + wherever such third-party notices normally appear. The contents + of the NOTICE file are for informational purposes only and + do not modify the License. You may add Your own attribution + notices within Derivative Works that You distribute, alongside + or as an addendum to the NOTICE text from the Work, provided + that such additional attribution notices cannot be construed + as modifying the License. + + You may add Your own copyright statement to Your modifications and + may provide additional or different license terms and conditions + for use, reproduction, or distribution of Your modifications, or + for any such Derivative Works as a whole, provided Your use, + reproduction, and distribution of the Work otherwise complies with + the conditions stated in this License. + + 5. Submission of Contributions. Unless You explicitly state otherwise, + any Contribution intentionally submitted for inclusion in the Work + by You to the Licensor shall be under the terms and conditions of + this License, without any additional terms or conditions. + Notwithstanding the above, nothing herein shall supersede or modify + the terms of any separate license agreement you may have executed + with Licensor regarding such Contributions. + + 6. Trademarks. This License does not grant permission to use the trade + names, trademarks, service marks, or product names of the Licensor, + except as required for reasonable and customary use in describing the + origin of the Work and reproducing the content of the NOTICE file. + + 7. Disclaimer of Warranty. Unless required by applicable law or + agreed to in writing, Licensor provides the Work (and each + Contributor provides its Contributions) on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or + implied, including, without limitation, any warranties or conditions + of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A + PARTICULAR PURPOSE. You are solely responsible for determining the + appropriateness of using or redistributing the Work and assume any + risks associated with Your exercise of permissions under this License. + + 8. Limitation of Liability. In no event and under no legal theory, + whether in tort (including negligence), contract, or otherwise, + unless required by applicable law (such as deliberate and grossly + negligent acts) or agreed to in writing, shall any Contributor be + liable to You for damages, including any direct, indirect, special, + incidental, or consequential damages of any character arising as a + result of this License or out of the use or inability to use the + Work (including but not limited to damages for loss of goodwill, + work stoppage, computer failure or malfunction, or any and all + other commercial damages or losses), even if such Contributor + has been advised of the possibility of such damages. + + 9. Accepting Warranty or Additional Liability. While redistributing + the Work or Derivative Works thereof, You may choose to offer, + and charge a fee for, acceptance of support, warranty, indemnity, + or other liability obligations and/or rights consistent with this + License. However, in accepting such obligations, You may act only + on Your own behalf and on Your sole responsibility, not on behalf + of any other Contributor, and only if You agree to indemnify, + defend, and hold each Contributor harmless for any liability + incurred by, or claims asserted against, such Contributor by reason + of your accepting any such warranty or additional liability. + + END OF TERMS AND CONDITIONS + + APPENDIX: How to apply the Apache License to your work. + + To apply the Apache License to your work, attach the following + boilerplate notice, with the fields enclosed by brackets "[]" + replaced with your own identifying information. (Don't include + the brackets!) The text should be enclosed in the appropriate + comment syntax for the file format. We also recommend that a + file or class name and description of purpose be included on the + same "printed page" as the copyright notice for easier + identification within third-party archives. + + Copyright [yyyy] [name of copyright owner] + + Licensed under the Apache License, Version 2.0 (the "License"); + you may not use this file except in compliance with the License. + You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + + Unless required by applicable law or agreed to in writing, software + distributed under the License is distributed on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + See the License for the specific language governing permissions and + limitations under the License. + +------------------------------------------------------------------------------------ +This product bundles various third-party components under other open source licenses. +This section summarizes those components and their licenses. See licenses/ +for text of these licenses. + +Apache Software Foundation License 2.0 +-------------------------------------- + +src/openai_api_protocols +``` + +--- + +## loglevel 1.9.2 + +- License: **MIT** +- Project: https://github.com/pimterry/loglevel +- Bundled transitively as WebLLM's only runtime dependency. + +``` +Copyright (c) 2013 Tim Perry + +Permission is hereby granted, free of charge, to any person +obtaining a copy of this software and associated documentation +files (the "Software"), to deal in the Software without +restriction, including without limitation the rights to use, +copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the +Software is furnished to do so, subject to the following +conditions: + +The above copyright notice and this permission notice shall be +included in all copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, +EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES +OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND +NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT +HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, +WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING +FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR +OTHER DEALINGS IN THE SOFTWARE. +``` diff --git a/WEBLLM-SURFACE.md b/WEBLLM-SURFACE.md new file mode 100644 index 0000000000000000000000000000000000000000..705378f39fea2c9ec405b099d220a3ae0231be62 --- /dev/null +++ b/WEBLLM-SURFACE.md @@ -0,0 +1,134 @@ +# WebLLM surface, and where the line is + +`@mlc-ai/web-llm` **0.2.84**, pinned exactly (`build.mjs` rewrites the bundle at anchors it verifies +first, and throws on a miss, so a minor bump breaks the build for every consumer). + +This file exists because it was needed and absent. Building "a superset of WebLLM" without a written +inventory of WebLLM meant reaching for `grep` one narrow question at a time, and three functions got +reimplemented that were exported all along — see "Removed", below. **Check this list before adding a +capability**, per AI.md's Reuse First. Re-check it on every dependency bump. + +**Every dependency listed here is asserted by +[test/webllm-contract.test.mjs](test/webllm-contract.test.mjs)** (`npm run contract`). Prose does not +fail a build; that file does. Add an assertion there whenever you add a row here. + +## Upgrading `@mlc-ai/web-llm` + +The pin is exact, and three independent guards catch three different kinds of drift. A bump is a +checklist that fails loudly at the right line, not an afternoon of `grep`. + +```sh +npm i @mlc-ai/web-llm@ +npm run build # 1. surface drift — patches verify before they rewrite +npm test # 2. semantic drift — the contract asserts every shape we depend on +npm run e2e # 3. behavioural drift — real GPU, byte-identical greedy generation +``` + +**1. Surface drift — `npm run build`.** [build/patches.mjs](build/patches.mjs) verifies all four +patch anchors against a fresh bundle *before* rewriting anything, so a bump reports every break at +once rather than the first over a half-patched file. A miss prints the anchor, the WebLLM version it +last held against, and the nearest lines ranked by identifier rarity; for a vanished identifier it +also names the closest surviving one by trigram overlap — an upstream rename of +`requiredMaxStorageBuffersPerShaderStage` shows up as an 85%-similar name with its line number. +Anchors match modulo whitespace and, where the literal string is generic (`compute.end();`), are +scoped to an enclosing function, so reflow and unrelated new code do not cause false failures. +`npm run verify-patches` runs just this check without rebuilding. Re-anchoring still needs a human to +*approve* the new location: edit `build/patches.mjs`, re-run, repeat until green. + +**2. Semantic drift — `npm test`.** [test/webllm-contract.test.mjs](test/webllm-contract.test.mjs) +(`npm run contract` alone, GPU-free) asserts every export, field, enum value and finish-reason +literal in the tables below, plus: + +- the ~30 undocumented tvmjs pipeline internals `multistep.js` reaches into — the list is *derived + from that file's own source* and pinned to `PIPELINE_CONTRACT`, so neither can fall behind the code +- `decode_tokens_per_s` still present in the bundle — `estimateSpeed()`'s measured path needs it +- `model_lib` still unguessable from base + id — the "require it for a remote source" rule in + ROADMAP.md's verb-consolidation section depends on this staying true + +A failure here is the dangerous kind: every symbol still resolves and behaviour changed anyway. Read +it before touching anything else. `test/patches.test.mjs` and `test/integration.test.mjs` also run +here and guard the patch machinery and cache-scope names. + +**3. Behavioural drift — `npm run e2e`.** Real Firefox, real GPU, real model, drag-and-drop +ingestion through the production `src/engine/` + `src/adapters/` paths. Asserts greedy generation is +byte-identical and the KV-reuse pipeline matches a forced re-prefill. The only guard that catches a +change where every name survives, every shape matches, and the tokens are still wrong. + +**Then, by hand:** + +- Bump the version line at the top of this file. +- Re-run the export list and diff it against the previous run: + ```sh + node -e 'import("./vendor/web-llm.js").then(m=>console.log(Object.keys(m).sort().join("\n")))' + ``` + A new export is a candidate for "available, not used yet". A *removed* one that appears in the + tables below should already have failed the contract test — if it did not, add the assertion. +- Walk the tables here for anything the diff or the test output touched. + +A large tvmjs refactor still needs a human to re-anchor and re-read. The goal is not an automatic +bump — it is that you find out what broke in seconds and know exactly where to look. + +## Native — we call it + +| Export | Used by | For | +| --- | --- | --- | +| `CreateWebWorkerMLCEngine` | `engine.js` | one engine per pool slot, each in its own realm | +| `WebWorkerMLCEngineHandler` | `engine-worker.js` | the worker side of that | +| `prebuiltAppConfig` | `engine.js` | 163 HuggingFace models, merged under our registry | +| `functionCallingModelIds` | `engine.js` | which prebuilt models do tool calling — `listAvailableModels()` tags them, `recommendModels({ needsToolCalling })` filters on them. Shipped in the list; not derivable from an id | +| `hasModelInCache` | `engine.cacheState()` | is a **remote/prebuilt** model on disk | +| `deleteModelAllInfoInCache` | `engine.evict()` | delete a **remote/prebuilt** model's bytes | +| `engine.chat.completions.create` | `pool.js` | the actual generation call | +| `engine.embeddings.create` | `pool.js` | `embed()` / `embedRaw()`. The OpenAI-shaped namespace (`Embeddings.prototype.create`), not the bare `MLCEngine.embedding()` — one forward pass, routed through the same scheduler as a completion | +| `usage.extra.decode_tokens_per_s` | `engine.#calibrate()` | measured throughput; calibrates projections | +| `finish_reason` (`stop`/`length`/`abort`) | `pool.js`, `chat.js` | truncation vs. natural stop vs. interrupt | +| `interruptGenerate()` | `pool.js` | cancellation and preemption | +| `unload()` → aborts `reloadController` | `engine.js` | **download cancellation is native**; no machinery of ours | +| `resetChat()` | `engine-worker.js` | forces re-prefill where paged KV cannot build | +| `ModelType` | `model-store.js` | `MODEL_TYPE`; a VLM must declare itself or images are refused | + +## Native — available, not used yet + +| Export | Why it matters | +| --- | --- | +| `CreateExtensionServiceWorkerMLCEngine` | **A native answer to the MV3 limit in AI.md.** Evaluate before hand-rolling a keep-alive. | +| `verifyIntegrity` / `isValidSRI` / `IntegrityError` | SRI checks on fetched artifacts — relevant once models are self-hosted. | +| `reload([a, b])` + routing on `request.model` | Multiple models in **one** engine. See "Overlapping by design". | +| `runtimeStatsText()` | Formatted prefill/decode stats. `usage.extra` is the structured form and is what we use. | +| `getGPUVendor()`, `getMaxStorageBufferBindingSize()` | Overlap `probeDevice()` but need a **loaded** engine, so they cannot answer "will this load". | +| `completion()`, `Completions` | Legacy non-chat completions. No caller has asked. | + +## Ours — nothing upstream does this + +| | Why it cannot be delegated | +| --- | --- | +| `pool.js` — priority bands, session supersession, opt-in preemption, one engine per task, demand-driven growth | WebLLM has no scheduler at all | +| `multistep.js` — K forward steps per GPU sync | 9.7 → 18.4 tok/s; not a WebLLM concept | +| `build.mjs` — compute-pass batching, storage-buffer clamp | 10.3 → 25.9 tok/s; patches *into* tvmjs | +| `ingest.js` — validate a folder, write Cache Storage directly | WebLLM only ever fetches; it cannot be handed bytes | +| `ModelStore.verify()` / `evictInjected()` | We wrote those artifacts and hold the only manifest. WebLLM's equivalents read `tensor-cache.json` to enumerate shards, so they break once *that* file is evicted — the exact case injection has to survive | +| `device.js` — `probeDevice`, `canRun`, `rankModels` | Preflight, before any download. WebLLM's device getters need a loaded engine | +| The engine-worker decode probe (`encodeMs`/`syncMs`) | Splits a burst into CPU-encode vs GPU-sync. Nothing upstream reports it, and the multi-step and pass-batching work is measured against it | +| `absolutize()` in `model-store.js` | WebLLM's `cleanModelUrl` ends in `new URL(url)` with no base, so a relative `/models/x/` throws deep in the loader. We resolve at registration | +| `errors.js`, the adapters, the wire protocol | Host and API concerns WebLLM has no opinion on | + +## Rejected — looked delegable, was not + +**Multi-model residency.** `reload(["a","b"])` holds N models in one engine and routes on +`request.model` — cheaper in memory, one realm instead of one worker per model. First presented here +as "not necessary in the way it was first presented", which was wrong: `reload()` calls `unload()` +unconditionally before loading, so it is **all-or-nothing** — adding a third model to `{A, B}` reloads +A and B too, ~51 s each on the 2B. There is no incremental residency upstream. Our `#pools: +Map` is the only thing that provides additive residency, which is a hard +requirement (switching models must not cost re-downloading the ones already up). See ROADMAP.md, +"Deferred", for the reopening condition. + +## Removed, having been found redundant + +| Was | Now | +| --- | --- | +| `cleanModelUrl()` reimplemented in `model-store.js` | deleted — only existed to support the two below | +| `ModelStore.cacheKeysFor()` | deleted — WebLLM derives the keys it fetched | +| `ModelStore.cacheState()` | `engine.cacheState()`, delegating to `hasModelInCache` for remote | +| `ModelStore.evict()` (all sources) | `ModelStore.evictInjected()` + `engine.evict()` routing to `deleteModelAllInfoInCache` | +| speed re-derived from the worker probe | `usage.extra.decode_tokens_per_s`, which every response already carried | diff --git a/package.json b/package.json new file mode 100644 index 0000000000000000000000000000000000000000..fde016356de1fe931a48c28f6ed84c644b7d2ac9 --- /dev/null +++ b/package.json @@ -0,0 +1,72 @@ +{ + "name": "everything-webgpu", + "version": "0.1.0", + "description": "Firefox WebExtension and embeddable library that runs local MLC/WebLLM models on WebGPU, with a priority scheduler, multi-step decoding, and models from HuggingFace, your own URL, or injected straight from disk.", + "scripts": { + "build": "node build.mjs", + "check": "node build.mjs --check-entries", + "test": "node --test test/webllm-contract.test.mjs test/patches.test.mjs test/readme.test.mjs test/api-doc.test.mjs test/examples.test.mjs test/license.test.mjs test/integration.test.mjs test/scheduler.test.mjs test/multistep.test.mjs test/environment.test.mjs test/sources.test.mjs test/prefetch.test.mjs test/embed.test.mjs test/recipes.test.mjs test/errors.test.mjs test/chat.test.mjs test/device.test.mjs test/manage.test.mjs", + "package": "npm run build && node build.mjs --zip", + "e2e": "node test/e2e/run.mjs", + "verify-consumer": "node test/consumer/verify.mjs", + "bench": "node test/e2e/bench.mjs", + "types": "tsc -p tsconfig.json", + "prepare": "node build.mjs", + "prepublishOnly": "npm run types && npm test", + "contract": "node --test test/webllm-contract.test.mjs", + "verify-patches": "node build.mjs --verify-patches" + }, + "author": "nyaaorick", + "license": "ISC", + "repository": { + "type": "git", + "url": "git+https://github.com/nyaaorick/everything-webgpu.git" + }, + "homepage": "https://github.com/nyaaorick/everything-webgpu#readme", + "bugs": "https://github.com/nyaaorick/everything-webgpu/issues", + "dependencies": { + "@mlc-ai/web-llm": "0.2.84" + }, + "devDependencies": { + "acorn": "^8.18.0", + "esbuild": "^0.28.2", + "typescript": "^5.9.2" + }, + "type": "module", + "types": "types/engine/index.d.ts", + "exports": { + ".": { + "types": "./types/engine/index.d.ts", + "default": "./src/engine/index.js" + }, + "./worker": "./src/engine/engine-worker.js", + "./vite": { + "types": "./types/vite.d.ts", + "default": "./src/vite.js" + }, + "./adapters/idb": { + "types": "./types/adapters/idb.d.ts", + "default": "./src/adapters/idb.js" + }, + "./adapters/memory": { + "types": "./types/adapters/memory.d.ts", + "default": "./src/adapters/memory.js" + }, + "./adapters/webext": { + "types": "./types/adapters/webext.d.ts", + "default": "./src/adapters/webext.js" + }, + "./adapters/protocol": { + "types": "./types/adapters/protocol.d.ts", + "default": "./src/adapters/protocol.js" + } + }, + "files": [ + "src/engine", + "src/adapters", + "src/vite.js", + "types", + "vendor/web-llm.js", + "THIRD-PARTY-NOTICES.md" + ] +} diff --git a/src/adapters/idb.js b/src/adapters/idb.js new file mode 100644 index 0000000000000000000000000000000000000000..cb7a2f0be87bc0d916244aff34c27147050f06c6 --- /dev/null +++ b/src/adapters/idb.js @@ -0,0 +1,84 @@ +/** + * IndexedDB StorageAdapter — the default for an ordinary page. + * + * IndexedDB rather than `localStorage` for two reasons that both matter here: + * `localStorage` is synchronous and blocks the thread that is also driving + * WebGPU, and it does not exist in a worker. The registry itself is small (one + * record per model, a few dozen cache URLs each), so this is a key/value store + * and nothing more. + * + * `persist()` is exported alongside because the page origin, unlike an + * extension with `unlimitedStorage`, holds a multi-GB model in *evictable* + * storage until persistence is granted. See ARCHIVE.md, "Zero-download". + */ +const DB_NAME = "everything-webgpu"; +const DB_VERSION = 1; +const STORE = "kv"; + +function open() { + return new Promise((resolve, reject) => { + const req = indexedDB.open(DB_NAME, DB_VERSION); + req.onupgradeneeded = () => req.result.createObjectStore(STORE); + req.onsuccess = () => resolve(req.result); + req.onerror = () => reject(req.error); + }); +} + +function run(db, mode, fn) { + return new Promise((resolve, reject) => { + const tx = db.transaction(STORE, mode); + const out = fn(tx.objectStore(STORE)); + tx.oncomplete = () => resolve(out); + tx.onerror = () => reject(tx.error); + tx.onabort = () => reject(tx.error); + }); +} + +const get1 = (store, key) => + new Promise((resolve, reject) => { + const req = store.get(key); + req.onsuccess = () => resolve(req.result); + req.onerror = () => reject(req.error); + }); + +/** @returns {Promise} */ +export async function indexedDBStorage() { + const db = await open(); + return { + async get(key) { + const keys = Array.isArray(key) ? key : [key]; + const out = {}; + await run(db, "readonly", async (store) => { + for (const k of keys) { + const value = await get1(store, k); + if (value !== undefined) out[k] = value; + } + }); + return out; + }, + async set(items) { + await run(db, "readwrite", (store) => { + for (const [k, v] of Object.entries(items)) store.put(v, k); + }); + }, + }; +} + +/** + * Ask the browser to stop treating this origin's storage as evictable. + * + * Firefox prompts the user (or grants silently for a site with the permission + * already); Chrome decides from engagement heuristics without a prompt. Either + * way a caller must handle `persisted: false` — the model still works, it can + * just be dropped under disk pressure, which `ModelStore.verify()` will catch + * on the next load. + * + * @returns {Promise<{persisted: boolean, quota?: number, usage?: number}>} + */ +export async function ensurePersistent() { + const storage = globalThis.navigator?.storage; + if (!storage?.persist) return { persisted: false }; + const persisted = (await storage.persisted?.()) || (await storage.persist()); + const estimate = await storage.estimate?.().catch(() => ({})); + return { persisted: Boolean(persisted), quota: estimate?.quota, usage: estimate?.usage }; +} diff --git a/src/adapters/memory.js b/src/adapters/memory.js new file mode 100644 index 0000000000000000000000000000000000000000..2e84a0ff4dde4433c670e65b3a3e136bc98331d5 --- /dev/null +++ b/src/adapters/memory.js @@ -0,0 +1,22 @@ +/** + * In-memory StorageAdapter. + * + * For tests, and for a host that genuinely wants the registry to die with the + * page. Note that the *weights* still live in Cache Storage and survive — only + * the registry entry pointing at them is lost, which `ModelStore.verify()` will + * then never be asked about. Do not use this in production for that reason. + */ + +/** @returns {import("../engine/model-store.js").StorageAdapter} */ +export function memoryStorage(initial = {}) { + const store = new Map(Object.entries(initial)); + return { + async get(key) { + const keys = Array.isArray(key) ? key : [key]; + return Object.fromEntries(keys.filter((k) => store.has(k)).map((k) => [k, store.get(k)])); + }, + async set(items) { + for (const [k, v] of Object.entries(items)) store.set(k, v); + }, + }; +} diff --git a/src/adapters/protocol.js b/src/adapters/protocol.js new file mode 100644 index 0000000000000000000000000000000000000000..1af68c72af5f7781e7bce00d087beacabf63fb6b --- /dev/null +++ b/src/adapters/protocol.js @@ -0,0 +1,57 @@ +/** + * Wire protocol — one adapter's vocabulary, not the engine's interface. + * + * A host that holds a `ScheduledEngine` calls its methods directly and never + * loads this file. This exists for the case where the engine and the caller are + * in different processes: a WebExtension hosting the engine in its background + * page, reached from a popup, an options page, or another extension. + * + * Two transports, one vocabulary: + * - `browser.runtime.sendMessage(EXTENSION_ID, req)` for request/response ops. + * - `browser.runtime.connect(EXTENSION_ID, { name: PORT_NAME })` for streaming. + * + * Every message carries `protocol` so a stray message from another sender fails + * fast instead of being half-interpreted. + * + * The engine is one GPU with a hard budget (see AI.md, "The 10 tok/s ceiling"), + * shared by every caller. So requests carry scheduling metadata, and the engine + * — not the caller — decides what runs when. + */ + +// Re-exported so a wire caller needs one import, not two. The definitions live +// with the scheduler, which is what actually acts on them. +export { ENGINE_STATE, PRIORITY, PRIORITY_ORDER } from "../engine/constants.js"; + +export const PROTOCOL = "everything-webgpu/v1"; +export const PORT_NAME = PROTOCOL; + +/** Request/response ops (runtime.sendMessage). */ +export const OP = { + STATUS: "status", + LIST_MODELS: "listModels", + LOAD: "load", + UNLOAD: "unload", + CHAT: "chat", + /** Many independent prompts at once; the engine fans them across the pool. */ + BATCH: "batch", + CANCEL: "cancel", + /** Retune a live engine (currently: `decodeSteps`) without reloading weights. */ + CONFIGURE: "configure", +}; + +/** Port ops (runtime.connect). */ +export const PORT_OP = { + CHAT_STREAM: "chat.stream", + BATCH_STREAM: "batch.stream", + ABORT: "abort", + SUBSCRIBE: "subscribe", + CHUNK: "chunk", + ITEM: "item", + DONE: "done", + ERROR: "error", + ENGINE_STATE: "engineState", +}; + +export function request(op, payload = {}) { + return { protocol: PROTOCOL, op, ...payload }; +} diff --git a/src/adapters/webext.js b/src/adapters/webext.js new file mode 100644 index 0000000000000000000000000000000000000000..b88a46e62120fbfca6e1e7dc2dd410b77d1cc1f0 --- /dev/null +++ b/src/adapters/webext.js @@ -0,0 +1,183 @@ +/** + * WebExtension adapter: the storage passthrough, and the message/port router. + * + * This is the whole extension-specific half of what used to be + * `src/background/background.js`. The engine no longer knows it exists — it is + * one way to reach a `ScheduledEngine`, and a page reaches the same object by + * calling its methods. + * + * The wire format is unchanged byte-for-byte from `everything-webgpu/v1`, so + * the popup, the manager page, the e2e devtest page and any external extension + * keep working with no edits. + */ +import { asEngineError } from "../engine/errors.js"; +import { OP, PORT_NAME, PORT_OP, PROTOCOL } from "./protocol.js"; + +/** + * `browser.storage.local` already *is* the StorageAdapter shape — `get(key)` + * and `set(obj)`. The two-method interface was chosen for that reason, so this + * is a passthrough rather than a translation layer. + * + * @returns {import("../engine/model-store.js").StorageAdapter} + */ +export function webExtensionStorage() { + return browser.storage.local; +} + +/** + * Wires a ScheduledEngine onto this extension's runtime messaging. + * + * @param {import("../engine/engine.js").ScheduledEngine} engine + * @returns {() => void} detach + */ +export function attachWebExtensionTransport(engine) { + const subscribers = new Set(); + + const unsubscribe = engine.subscribe((state) => { + broadcast({ protocol: PROTOCOL, op: PORT_OP.ENGINE_STATE, state }); + }); + + function broadcast(msg) { + for (const port of subscribers) { + try { + port.postMessage(msg); + } catch { + subscribers.delete(port); + } + } + } + + // ------------------------------------------------------------- routing ---- + + async function handle(msg) { + if (!msg || msg.protocol !== PROTOCOL) { + throw new Error(`Expected protocol "${PROTOCOL}", got "${msg && msg.protocol}".`); + } + switch (msg.op) { + case OP.STATUS: + return { ok: true, state: engine.state, webgpu: engine.hasWebGPU }; + case OP.LIST_MODELS: + return { ok: true, models: await engine.listModels(), state: engine.state }; + case OP.LOAD: + return { ok: true, state: await engine.load(msg.modelId) }; + case OP.UNLOAD: + return { ok: true, state: await engine.unload() }; + case OP.CHAT: + return { ok: true, ...(await engine.complete(msg)) }; + case OP.BATCH: + return { ok: true, results: await engine.batch(msg) }; + case OP.CANCEL: + return { ok: true, cancelled: engine.cancel(msg.id ?? msg.session) }; + case OP.CONFIGURE: + return { ok: true, ...(await engine.configure(msg)) }; + default: + throw new Error(`Unknown op "${msg.op}".`); + } + } + + /** + * `error` stays a plain string so every existing caller keeps working; `code` + * and `detail` are added beside it. Drop the string once consumers have + * moved — the wire protocol is versioned (`everything-webgpu/v1`) and this is + * an addition, not a break. + */ + const respond = (msg) => + handle(msg).catch((err) => { + const e = asEngineError(err); + return { ok: false, error: e.message, code: e.code, ...(e.detail ? { detail: e.detail } : {}) }; + }); + + /** + * The allowlist is a property of *this transport*, not of the engine: it + * answers "which other extension may send me messages", a question that has + * no meaning for a page that already holds the object. + */ + async function denyExternal(sender) { + const { allowedExternalIds } = await engine.store.getSettings(); + const id = sender?.id; + if (allowedExternalIds.length === 0 || allowedExternalIds.includes(id)) return null; + return { ok: false, error: `Extension "${id}" is not on this engine's allowlist.` }; + } + + function attachPort(port) { + if (port.name !== PORT_NAME) return; + subscribers.add(port); + port.onDisconnect.addListener(() => subscribers.delete(port)); + port.postMessage({ protocol: PROTOCOL, op: PORT_OP.ENGINE_STATE, state: engine.state }); + + const send = (msg) => { + try { + port.postMessage(msg); + } catch { + /* port closed mid-stream */ + } + }; + + port.onMessage.addListener(async (msg) => { + const id = msg?.id; + try { + if (!msg || msg.protocol !== PROTOCOL) throw new Error(`Expected protocol "${PROTOCOL}".`); + switch (msg.op) { + case PORT_OP.SUBSCRIBE: + return send({ protocol: PROTOCOL, op: PORT_OP.ENGINE_STATE, state: engine.state }); + + case PORT_OP.CHAT_STREAM: { + const result = await engine.complete(msg, (delta) => + send({ protocol: PROTOCOL, op: PORT_OP.CHUNK, id, delta }), + ); + return send({ protocol: PROTOCOL, op: PORT_OP.DONE, id, ...result }); + } + + case PORT_OP.BATCH_STREAM: { + const results = await engine.batch(msg, (item) => + send({ protocol: PROTOCOL, op: PORT_OP.ITEM, id, ...item }), + ); + return send({ protocol: PROTOCOL, op: PORT_OP.DONE, id, results }); + } + + case PORT_OP.ABORT: + return void engine.cancel(msg.session ?? id); + + default: + return send({ protocol: PROTOCOL, op: msg.op, id, ...(await respond(msg)) }); + } + } catch (err) { + const e = asEngineError(err); + send({ + protocol: PROTOCOL, + op: PORT_OP.ERROR, + id, + error: e.message, + code: e.code, + ...(e.detail ? { detail: e.detail } : {}), + }); + } + }); + } + + const onMessage = (msg) => respond(msg); + const onMessageExternal = async (msg, sender) => (await denyExternal(sender)) ?? respond(msg); + const onConnectExternal = async (port) => { + const denial = await denyExternal(port.sender); + if (denial) { + port.postMessage({ protocol: PROTOCOL, op: PORT_OP.ERROR, error: denial.error }); + port.disconnect(); + return; + } + attachPort(port); + }; + + browser.runtime.onMessage.addListener(onMessage); + browser.runtime.onMessageExternal.addListener(onMessageExternal); + browser.runtime.onConnect.addListener(attachPort); + browser.runtime.onConnectExternal.addListener(onConnectExternal); + + return () => { + browser.runtime.onMessage.removeListener(onMessage); + browser.runtime.onMessageExternal.removeListener(onMessageExternal); + browser.runtime.onConnect.removeListener(attachPort); + browser.runtime.onConnectExternal.removeListener(onConnectExternal); + unsubscribe(); + subscribers.clear(); + }; +} diff --git a/src/background/background.html b/src/background/background.html new file mode 100644 index 0000000000000000000000000000000000000000..412379d0a79dbab5dcb136500fc742d0f19ca1c7 --- /dev/null +++ b/src/background/background.html @@ -0,0 +1,4 @@ + + +Everything WebGPU — engine host + diff --git a/src/background/background.js b/src/background/background.js new file mode 100644 index 0000000000000000000000000000000000000000..0036b420b974ffc48389baef8512ab52e0256d8f --- /dev/null +++ b/src/background/background.js @@ -0,0 +1,22 @@ +/** + * The extension's engine host: build a ScheduledEngine, attach the transport. + * + * Everything that used to be in this file is now either engine (`src/engine/`) + * or transport (`src/adapters/webext.js`). What remains is the two lines that + * are genuinely specific to *this* host — which storage to use, and which + * transport to speak — and that is the whole shape a consumer of the library + * writes for their own app. + * + * The engine lives in the MV2 persistent background page because that is a real + * document on the extension origin, so it has both `navigator.gpu` and the same + * Cache Storage the manager page writes to. The model stays resident in VRAM + * across popup opens and across calls from other extensions. + */ +import { ScheduledEngine } from "../engine/index.js"; +import { attachWebExtensionTransport, webExtensionStorage } from "../adapters/webext.js"; + +export const engine = new ScheduledEngine({ store: webExtensionStorage() }); + +attachWebExtensionTransport(engine); + +console.info("[Everything WebGPU] engine host ready; WebGPU present:", engine.hasWebGPU); diff --git a/src/engine/chat.js b/src/engine/chat.js new file mode 100644 index 0000000000000000000000000000000000000000..86b355c9a461caf9da1083d0caf919573b31c5f0 --- /dev/null +++ b/src/engine/chat.js @@ -0,0 +1,142 @@ +/** + * `engine.chat.completions.create()` — the WebLLM-shaped facade. + * + * The point of this file is that migrating off `@mlc-ai/web-llm` costs one + * line. WebLLM is OpenAI-shaped, `buildParams()` already forwards `messages`, + * `temperature`, `max_tokens`, `response_format` and `extra_body` untouched, and + * everything this engine adds — the scheduler, multi-step decoding, the two + * build patches, the three model sources — sits behind that same call rather + * than beside it. So the facade is as thin as it can be: streamed chunks are + * WebLLM's own objects, passed through untouched, and only the non-streaming + * response is assembled here. + * + * `created` is therefore WebLLM's stable per-response value in milliseconds + * (OpenAI uses seconds; WebLLM does not, and the drop-in target is WebLLM), and + * the finish reasons are its own: `"stop" | "length" | "abort" | "tool_calls"`. + * + * `complete()` and `batch()` remain the direct API. They expose `cancelled` and + * `preempted` as first-class outcomes, which the OpenAI shape has no room for — + * both collapse to `finish_reason: "abort"` here, with the flags carried + * alongside for a caller that cares which happened. + */ +import { ERROR, EngineError } from "./errors.js"; + +/** @param {import("./engine.js").ScheduledEngine} engine */ +export function chatFacade(engine) { + return { + completions: { + /** + * @param {import("./engine.js").CompletionRequest & { + * stream?: boolean, stream_options?: {include_usage?: boolean} }} req + * @returns {Promise>} a completion, or a + * stream of chunks when `stream` is set — the same two shapes WebLLM + * returns, so `await`ing then `for await`ing works unchanged. + */ + async create(req) { + if (!Array.isArray(req?.messages) || req.messages.length === 0) { + throw new EngineError(ERROR.BAD_REQUEST, "`messages` must be a non-empty array."); + } + return req.stream ? streamCompletion(engine, req) : oneCompletion(engine, req); + }, + }, + }; +} + +const envelope = (engine, id, object) => ({ + id, + object, + created: Date.now(), + model: engine.state.modelId, +}); + +async function oneCompletion(engine, req) { + const id = req.id ?? crypto.randomUUID(); + const result = await engine.completeRaw(req); + return { + ...envelope(engine, id, "chat.completion"), + choices: [ + { + index: 0, + // `content: null` beside tool_calls is WebLLM's own shape, not "". + message: result.toolCalls + ? { role: "assistant", content: null, tool_calls: result.toolCalls } + : { role: "assistant", content: result.text }, + finish_reason: result.finishReason ?? "stop", + logprobs: null, + }, + ], + usage: result.usage, + ...(result.cancelled ? { cancelled: true } : {}), + ...(result.preempted ? { preempted: true } : {}), + }; +} + +/** + * Bridges the raw chunk callback to an async iterator. + * + * Chunks pass through **verbatim**. WebLLM's are already compliant OpenAI + * envelopes carrying `id`, `created`, `model`, `logprobs`, + * `system_fingerprint` and the terminal `tool_calls`; the previous version + * rebuilt them from a bare string and lost all of that. + * + * Chunks are queued rather than awaited, because the engine must not be made to + * wait on a slow consumer: a stalled `for await` would hold a pool slot, and a + * pool slot is the scarce resource the whole scheduler exists to allocate. The + * queue is bounded in practice by `max_tokens`. + */ +async function streamCompletion(engine, req) { + const id = req.id ?? crypto.randomUUID(); + + /** @type {object[]} */ + const pending = []; + let wake = null; + /** @type {{result?: object, error?: unknown} | null} */ + let settled = null; + const ping = () => { + const w = wake; + wake = null; + w?.(); + }; + + engine + .completeRaw(req, (chunk) => { + pending.push(chunk); + ping(); + }) + .then( + (result) => { + settled = { result }; + ping(); + }, + (error) => { + settled = { error }; + ping(); + }, + ); + + return (async function* () { + for (;;) { + // Drain before checking `settled`, so the last chunks are never dropped + // by a generation that finished while they sat in the queue. + while (pending.length) yield pending.shift(); + if (settled) break; + await new Promise((resolve) => (wake = resolve)); + } + + if (settled.error) throw settled.error; + const result = settled.result; + + // Nothing is synthesized on the normal path: WebLLM emits its own terminal + // finish_reason chunk, and its own usage chunk when `include_usage` is set. + // An interrupted generation is the exception — the stream simply stops, so + // a consumer would otherwise never learn why. + if (result.cancelled || result.preempted) { + yield { + ...envelope(engine, id, "chat.completion.chunk"), + choices: [{ index: 0, delta: {}, finish_reason: "abort" }], + ...(result.cancelled ? { cancelled: true } : {}), + ...(result.preempted ? { preempted: true } : {}), + }; + } + })(); +} diff --git a/src/engine/constants.js b/src/engine/constants.js new file mode 100644 index 0000000000000000000000000000000000000000..2ad81171227fc4d43fb7b4aaa7eecf77d4f8b27b --- /dev/null +++ b/src/engine/constants.js @@ -0,0 +1,68 @@ +/** + * Engine vocabulary — the words the scheduler itself is defined in. + * + * Deliberately transport-free. A caller holding a `ScheduledEngine` uses these + * directly; a caller reaching the engine over a message port gets them + * re-exported from `src/adapters/protocol.js` alongside the wire ops. Keeping + * the split visible is the point: the wire protocol is one adapter, not the + * interface. + */ + +/** + * Scheduling bands, highest first. `interactive` may preempt a running job that + * opted into `preemptible`; nothing else ever interrupts work in flight. + */ +export const PRIORITY = { + /** Ghost-text and anything else a human is waiting on keystroke-by-keystroke. */ + INTERACTIVE: "interactive", + /** Default. Translation, one-shot answers. */ + NORMAL: "normal", + /** Nobody is watching: reformatting, batch cleanup. Pair with `preemptible`. */ + BACKGROUND: "background", +}; + +export const PRIORITY_ORDER = [PRIORITY.INTERACTIVE, PRIORITY.NORMAL, PRIORITY.BACKGROUND]; + +/** + * How far `unload()` goes. + * + * Two levels rather than two verbs because they are the same intention at + * different depths — "I am done with this model" — and the caller should not + * have to know that freeing VRAM and freeing disk are different subsystems. + * Forgetting the model *entirely* is `remove()`, which stays its own verb + * because it is the one that cannot be undone without re-supplying the source. + */ +export const UNLOAD_LEVEL = { + /** Default. Free VRAM, keep the cached bytes — so reloading costs no network. */ + VRAM: "vram", + /** Also delete the cached bytes. The registry entry survives, so the model is still known. */ + CACHE: "cache", +}; + +/** + * What a queued job asks its engine to do. + * + * One pool, not two: priority, session supersession, preemption and + * one-task-one-engine are identical for both, and the only thing that differs + * is the call at the far end. A second pool would have duplicated the + * scheduler to change one line. + */ +export const JOB_KIND = { + CHAT: "chat", + EMBEDDING: "embedding", +}; + +/** Engine lifecycle states. */ +export const ENGINE_STATE = { + IDLE: "idle", + LOADING: "loading", + READY: "ready", + ERROR: "error", +}; + +/** + * Out-of-band message kind for configuring an engine worker before WebLLM's own + * handshake starts. Distinct from every WebLLM `kind`, so the worker can route + * on it without parsing the rest. + */ +export const WORKER_CONFIGURE = "everything-webgpu/configure"; diff --git a/src/engine/create.js b/src/engine/create.js new file mode 100644 index 0000000000000000000000000000000000000000..487276367fb6184479ba3e8168bdf4bf44aeed8c --- /dev/null +++ b/src/engine/create.js @@ -0,0 +1,138 @@ +/** + * `CreateScheduledEngine` — the one-line migration off WebLLM. + * + * ```js + * -import { CreateMLCEngine } from "@mlc-ai/web-llm"; + * -const engine = await CreateMLCEngine(modelId, { initProgressCallback }); + * +import { CreateScheduledEngine } from "everything-webgpu"; + * +const engine = await CreateScheduledEngine(modelId, { initProgressCallback }); + * ``` + * + * Everything after that line is unchanged: `engine.chat.completions.create()` + * takes and returns the same shapes. What the swap buys is the scheduler + * (`session`, `priority`, `task`, `preemptible` on any request), multi-step + * decoding, the two build-time patches, and models from disk as well as the + * network. + * + * The signature mirrors `CreateMLCEngine`'s so nothing else has to move — + * including `initProgressCallback`, which is why this is a function rather than + * documentation telling people to construct a `ScheduledEngine` themselves. + */ +import { ScheduledEngine } from "./engine.js"; +import { ERROR, EngineError } from "./errors.js"; +import { ModelStore } from "./model-store.js"; + +/** + * @param {string} [modelId] loaded before returning, as `CreateMLCEngine` does. + * Omit to get an engine that loads later. + * @param {object} [opts] + * @param {ModelStore | import("./model-store.js").StorageAdapter} [opts.store] + * defaults to IndexedDB. Pass one explicitly in a worker or a test. + * @param {(report: {text: string, progress: number}) => void} [opts.initProgressCallback] + * @param {string | URL} [opts.workerUrl] + * @param {() => Promise} [opts.loadWebLLM] + * @param {boolean} [opts.prebuilt] + * @returns {Promise} + */ +export async function CreateScheduledEngine(modelId, opts = {}) { + const { store, initProgressCallback, ...rest } = opts; + const engine = new ScheduledEngine({ ...rest, store: store ?? (await defaultStore()) }); + if (!modelId) return engine; + + // Three states, not two. `undefined` means the caller did not choose, and the + // right default for a call that blocks for minutes on a multi-hundred-MB + // download is to say so: the first run of the three-line example otherwise + // looks exactly like a hung process, which is the worst possible first + // impression and the one thing no amount of documentation undoes. `null` is + // how you ask for silence. + const report = + initProgressCallback === undefined ? defaultProgressReporter(modelId) : initProgressCallback; + + if (!report) { + await engine.load(modelId); + return engine; + } + + // The engine reports progress through its own lifecycle stream; WebLLM + // reports it through a callback. Forward only the load reports, so a + // caller's callback fires exactly when WebLLM's would have. + const stop = engine.subscribe((state) => { + if (state.progress) report(state.progress); + }); + try { + await engine.load(modelId); + } finally { + stop(); + report.done?.(); + } + return engine; +} + +/** + * What you get when you say nothing: a throttled line on the console. + * + * Deliberately not the engine's own behaviour — `new ScheduledEngine()` stays + * silent, because a library core that logs is wrong inside a worker, an + * extension background page or a test. This is the getting-started facade, and + * the person calling it has not yet decided how to render anything. + * + * Throttled to one line a second, because WebLLM's callback fires per shard and + * a 58-shard model would otherwise bury the console. The 100% report is never + * dropped, so the last line always reads as finished rather than as 97%. + */ +function defaultProgressReporter(modelId) { + if (typeof console === "undefined") return null; + + const startedAt = Date.now(); + let announced = false; + let lastAt = 0; + let sawProgress = false; + + const emit = ({ text, progress }) => { + sawProgress = true; + if (!announced) { + announced = true; + console.info( + `[everything-webgpu] loading ${modelId} — the first run downloads the weights and can take ` + + "minutes; later loads read the cache and need no network. " + + "Pass initProgressCallback to render this yourself, or null to silence it.", + ); + } + const now = Date.now(); + const finished = (progress ?? 0) >= 1; + if (!finished && now - lastAt < 1000) return; + lastAt = now; + const pct = `${Math.round((progress ?? 0) * 100)}%`.padStart(4); + console.info(`[everything-webgpu] ${pct} ${text ?? ""}`.trimEnd()); + }; + + // Called once when the load settles, so the common case ends on a line that + // says it worked rather than trailing off mid-progress. + emit.done = () => { + if (!sawProgress) return; + console.info( + `[everything-webgpu] ${modelId} ready in ${((Date.now() - startedAt) / 1000).toFixed(1)}s.`, + ); + }; + return emit; +} + +/** + * IndexedDB, because a registry that dies with the page would strand the + * weights: the bytes stay in Cache Storage but nothing remembers they are + * there. Anywhere IndexedDB is missing, the caller has to say what to use — + * guessing would produce exactly that stranding, silently. + */ +async function defaultStore() { + if (typeof indexedDB === "undefined") { + throw new EngineError( + ERROR.BAD_REQUEST, + "No IndexedDB in this context, so there is no default store. " + + "Pass `store` — `memoryStorage()` for a test, or your own StorageAdapter.", + ); + } + // Imported lazily so the engine core carries no static dependency on an + // adapter, and a host that brings its own store never bundles this one. + const { indexedDBStorage } = await import("../adapters/idb.js"); + return new ModelStore(await indexedDBStorage()); +} diff --git a/src/engine/device.js b/src/engine/device.js new file mode 100644 index 0000000000000000000000000000000000000000..725f46a64fe19dfe9b9820369616b32c6e7d87c9 --- /dev/null +++ b/src/engine/device.js @@ -0,0 +1,279 @@ +/** + * "Will this model run here, and if not, why not" — answered before a byte is + * fetched. + * + * Every input is readable up front: WebGPU presence, whether an adapter can be + * had at all, `shader-f16`, five adapter limits, and the storage quota. So the + * expensive failure — download a gigabyte, then fail at pipeline creation — is + * avoidable, and this file exists to avoid it. + * + * It is also where this project's hard-won platform knowledge is written down + * as code rather than prose. Each rule below is a failure that actually + * happened here: the blocklisted GPU, Firefox's 9-storage-buffer cap that kills + * paged prefill, `q4f16_1` on a device without f16. + * + * What it deliberately does **not** claim: how much VRAM is free. Firefox + * implements neither `navigator.deviceMemory` nor `performance.memory`, and + * `storage.estimate()` measures disk quota, not memory. Nothing reports free + * memory to a web page. So VRAM headroom is reported as a warning with the + * numbers that *are* known, never as a confident blocker — the same reason the + * pool probes for a second engine rather than predicting one (AI.md, "The pool + * grows, it is not sized"). + */ + +/** WebLLM's ModelType enum; see model-store.js. */ +const MODEL_TYPE_LLM = 0; +const MODEL_TYPE_VLM = 2; + +/** Bindings `batch_prefill_paged_kv_kernel` needs; below this, no KV reuse. */ +const PAGED_PREFILL_STORAGE_BUFFERS = 10; + +/** + * @typedef {object} DeviceProbe + * @property {boolean} webgpu + * @property {string} [reason] why WebGPU is unusable, when it is + * @property {object} [adapter] vendor / architecture / device, where exposed + * @property {{shaderF16: boolean}} [features] + * @property {object} [limits] + * @property {boolean} [kvReuse] whether cross-turn KV reuse can be used here + * @property {{quota?: number, usage?: number, persisted?: boolean}} storage + * @property {number} [deviceMemoryGB] Chrome only; absent is not "small" + */ + +/** + * Reads what this machine will admit to. Never throws: an unusable device is a + * result, not an error — the caller's job is to explain it, not to crash. + * + * @returns {Promise} + */ +export async function probeDevice() { + const storage = await probeStorage(); + const gpu = globalThis.navigator?.gpu; + + if (!gpu) { + return { + webgpu: false, + reason: + "navigator.gpu is absent. On Firefox set dom.webgpu.enabled=true in about:config and restart; " + + "on any browser, a non-secure context (plain http) also hides it.", + storage, + }; + } + + // A `navigator.gpu` that is present but incomplete — a polyfill, a shim, a + // partially-enabled build — must be a result like any other, not a TypeError + // out of a function documented never to throw. + const adapter = + typeof gpu.requestAdapter === "function" ? await gpu.requestAdapter().catch(() => null) : null; + if (!adapter) { + return { + webgpu: false, + reason: + "WebGPU is present but no adapter was granted — usually a blocklisted GPU. On Firefox try " + + "gfx.webgpu.ignore-blocklist=true in about:config.", + storage, + }; + } + + const limits = {}; + for (const key of [ + "maxBufferSize", + "maxStorageBufferBindingSize", + "maxStorageBuffersPerShaderStage", + "maxComputeInvocationsPerWorkgroup", + "maxComputeWorkgroupStorageSize", + ]) { + limits[key] = adapter.limits?.[key]; + } + + return { + webgpu: true, + adapter: await adapterInfo(adapter), + features: { shaderF16: Boolean(adapter.features?.has?.("shader-f16")) }, + limits, + // Below 10 bindings the paged-prefill pipeline cannot be built, so every + // turn re-prefills the whole history. See engine-worker.js. + kvReuse: (limits.maxStorageBuffersPerShaderStage ?? 0) >= PAGED_PREFILL_STORAGE_BUFFERS, + storage, + ...(globalThis.navigator?.deviceMemory ? { deviceMemoryGB: navigator.deviceMemory } : {}), + }; +} + +/** `adapter.info` is the current spec; `requestAdapterInfo()` was the old one. */ +/** + * What the adapter will admit about itself, with blanks dropped. + * + * Firefox 154 exposes `adapter.info` but fills every field with `""`, so the + * naive shape is an object that *looks* populated and renders as "GPU: ". A + * caller cannot tell that from a real answer without checking each string, so + * empty fields are omitted and a browser that says nothing yields `{}` — the + * same thing the no-info path already returns. + */ +async function adapterInfo(adapter) { + const info = adapter.info ?? (await adapter.requestAdapterInfo?.().catch(() => null)); + if (!info) return {}; + const { vendor, architecture, device, description } = info; + return Object.fromEntries( + Object.entries({ vendor, architecture, device, description }).filter(([, v]) => v), + ); +} + +async function probeStorage() { + const s = globalThis.navigator?.storage; + if (!s) return {}; + const [estimate, persisted] = await Promise.all([ + s.estimate?.().catch(() => ({})) ?? {}, + s.persisted?.().catch(() => undefined) ?? undefined, + ]); + return { quota: estimate?.quota, usage: estimate?.usage, persisted }; +} + +/** + * Whether a model can run on a probed device. + * + * `blockers` mean it will not work; `warnings` mean it will work worse, or + * might not fit. The split matters: a caller should refuse to start on a + * blocker and merely inform on a warning, and conflating the two is how you end + * up either crashing or refusing to run something that would have been fine. + * + * @param {{model_id: string, vram_required_MB?: number, sizeBytes?: number}} model + * @param {DeviceProbe} probe + * @returns {{ok: boolean, blockers: Array<{code: string, message: string}>, + * warnings: Array<{code: string, message: string}>}} + */ +export function canRun(model, probe) { + const blockers = []; + const warnings = []; + const say = (list, code, message) => list.push({ code, message }); + + if (!probe?.webgpu) { + say(blockers, "NO_WEBGPU", probe?.reason ?? "WebGPU is unavailable."); + return { ok: false, blockers, warnings }; + } + + // The quantisation is encoded in the model id by MLC convention + // (`…-q4f16_1-MLC`). It is a naming convention, not a manifest field, so this + // is a heuristic — but a wrong guess only costs a spurious warning, while not + // checking costs a gigabyte downloaded before a pipeline fails. + if (/f16/.test(model?.model_id ?? "") && !probe.features?.shaderF16) { + say( + blockers, + "NO_SHADER_F16", + `"${model.model_id}" is an f16 build and this device has no \`shader-f16\` feature. ` + + "Choose a q4f32_1 or q0f32 variant — they are larger and slower, but they will run.", + ); + } + + if (!probe.kvReuse) { + say( + warnings, + "NO_KV_REUSE", + `This device allows ${probe.limits?.maxStorageBuffersPerShaderStage} storage buffers per shader ` + + `stage; paged prefill needs ${PAGED_PREFILL_STORAGE_BUFFERS}. Every turn re-prefills the whole ` + + "history, so long conversations get slow first tokens.", + ); + } + + const vramMB = model?.vram_required_MB; + if (vramMB && probe.deviceMemoryGB && vramMB > probe.deviceMemoryGB * 1024 * 0.5) { + say( + warnings, + "TIGHT_MEMORY", + `The model wants ~${Math.round(vramMB)} MB and this device reports ${probe.deviceMemoryGB} GB of ` + + "memory. It may load and then run against swap, or fail to load a second engine.", + ); + } + + const bytes = model?.sizeBytes ?? (vramMB ? vramMB * 1024 * 1024 : 0); + const free = probe.storage?.quota != null ? probe.storage.quota - (probe.storage.usage ?? 0) : null; + if (bytes && free != null && bytes > free) { + say( + blockers, + "INSUFFICIENT_STORAGE", + `The model needs ~${Math.round(bytes / 1e6)} MB cached and only ~${Math.round(free / 1e6)} MB of ` + + "quota is free.", + ); + } + + if (bytes && probe.storage?.persisted === false) { + say( + warnings, + "NOT_PERSISTED", + "Storage is not persisted, so the browser may evict the model under disk pressure. " + + "Call ensurePersistent() to ask for it.", + ); + } + + return { ok: blockers.length === 0, blockers, warnings }; +} + +/** + * Rank a model list by what this device can actually run. + * + * The prebuilt list spans 239 MB to 31 GB, so "which model should I use" is the + * first question a developer has and the one they have least basis to answer. + * Runnable models come first, then fewest warnings; unrunnable ones are kept at + * the end carrying their reason rather than silently dropped, because "why + * can't I use that one" is the next question. + * + * **`prefer` is a real choice, not a default worth hiding.** Decode here is + * memory-bandwidth-bound — time per token scales with weight bytes (AI.md, + * "Why not llama.cpp/Ollama-class"), so the largest model that fits is also the + * slowest thing that fits. `"quality"` picks the biggest, `"speed"` the + * smallest. Neither is right for everyone, which is why it is a parameter. + * + * Vision models are excluded from a text ranking rather than merely deprioritised: + * a VLM answers text prompts perfectly well, but at several times the download + * for no benefit, so recommending one to a caller who did not ask is bad advice. + * + * @param {Array} models `model_list` entries or registry records + * @param {{probe: DeviceProbe, maxVramMB?: number, needsVision?: boolean, + * prefer?: "quality" | "speed"}} opts + */ +export function rankModels(models, { probe, maxVramMB, needsVision = false, prefer = "quality" } = {}) { + const wanted = needsVision ? MODEL_TYPE_VLM : MODEL_TYPE_LLM; + return models + .filter((m) => (m.model_type ?? MODEL_TYPE_LLM) === wanted) + .filter((m) => !maxVramMB || !m.vram_required_MB || m.vram_required_MB <= maxVramMB) + .map((m) => ({ model: m, ...canRun(m, probe) })) + .sort((a, b) => { + if (a.ok !== b.ok) return a.ok ? -1 : 1; + if (a.warnings.length !== b.warnings.length) return a.warnings.length - b.warnings.length; + const sizeA = a.model.vram_required_MB ?? 0; + const sizeB = b.model.vram_required_MB ?? 0; + return prefer === "speed" ? sizeA - sizeB : sizeB - sizeA; + }); +} + +/** + * Decode throughput is memory bandwidth divided by weight bytes. + * + * This project measured the whole chain: decode reaches ~16 GB/s of the M4's + * ~120 GB/s, and the 1.06 GB build runs 16.6–18.1 tok/s — which is that + * quotient. So a projection needs one number, the *achieved* bandwidth, and + * everything else follows from model size. + * + * The constant below is that machine's figure and is only a starting point. The + * moment this engine has decoded anything it knows the real number for the + * machine it is on, and `ScheduledEngine.estimateSpeed()` switches to it — so + * this is a cold-start default, not a claim about anyone's hardware. + */ +export const REFERENCE_DECODE_BYTES_PER_SECOND = 17e9; +export const REFERENCE_DEVICE = "M4 MacBook Air (16 GB), Firefox"; + +/** + * @param {number} modelBytes + * @param {number} [bytesPerSecond] this machine's measured rate, when known + * @returns {{tokensPerSecond: number, basis: "measured" | "extrapolated", + * modelBytes: number, bytesPerSecond: number, reference?: string}} + */ +export function projectSpeed(modelBytes, bytesPerSecond) { + const rate = bytesPerSecond || REFERENCE_DECODE_BYTES_PER_SECOND; + return { + tokensPerSecond: modelBytes > 0 ? rate / modelBytes : 0, + basis: bytesPerSecond ? "measured" : "extrapolated", + modelBytes, + bytesPerSecond: rate, + ...(bytesPerSecond ? {} : { reference: REFERENCE_DEVICE }), + }; +} diff --git a/src/engine/engine-worker.js b/src/engine/engine-worker.js new file mode 100644 index 0000000000000000000000000000000000000000..57409c50fa3a64e501cde0bfe5018d99bb4b9bda --- /dev/null +++ b/src/engine/engine-worker.js @@ -0,0 +1,142 @@ +/** + * One pool slot's engine, in its own JS realm. + * + * Several MLCEngines cannot share a realm. Isolated by running the same e2e + * three ways against Qwen3.5-0.8B: + * + * 1 engine, background page -> passes + * 2 engines, background page -> both load, the first generates + * fine, the second's first + * generation fails with + * "Expected null or instance of + * VectorInt, got an instance of + * VectorInt" + * 2 engines, one worker each -> passes + * + * So the trigger is a second engine generating in the same realm, not the pool + * or the engine count as such. That error is embind reporting a type-registry + * mismatch, and the bundle does carry module-scoped emscripten state + * (`var Module`, `var __wasmLib`) shared by every instance, which fits — but + * the fix rests on the isolation above, not on having traced the registry. + * + * Workers are viable because Firefox exposes WebGPU to dedicated workers and + * the 100 ms completion tick is shared across them, so the concurrency win + * survives the move off the main thread (measured: 4 workers, 36.3 syncs/s). + * + * The realm is also where the decode loop lives, so it is where multi-step + * decoding is installed — the background page only ever holds a proxy. + */ +import { WebWorkerMLCEngineHandler } from "../../vendor/web-llm.js"; +import { WORKER_CONFIGURE } from "./constants.js"; +import { DEFAULT_DECODE_STEPS, installMultiStepDecoding } from "./multistep.js"; + +const handler = new WebWorkerMLCEngineHandler(); + +/** + * Running totals for the decode probe (AI.md, "Where the 46 ms goes"). + * + * `encodeMs` is content-process CPU — command encoding, `createBindGroup`, IPC. + * `syncMs` is GPU execution plus the wait for Firefox's 100 ms poll tick. They + * are measured on either side of the burst's single `await`, so together they + * partition the decode budget with nothing unaccounted for. + */ +const stats = { + bursts: 0, + tokens: 0, + encodeMs: 0, + syncMs: 0, + dispatches: 0, + forwardDispatches: 0, + flushes: 0, +}; +const resetStats = () => Object.keys(stats).forEach((k) => (stats[k] = 0)); + +/** + * Set when a pipeline fails the multi-step contract, and never cleared — a + * retune resets the measurement window, not the fact that the fast path is off. + */ +let multiStepOff = null; + +// No `kind` field: WebLLM's client handler ignores messages it cannot classify +// instead of throwing UnknownMessageKindError. +const postStats = () => + postMessage({ ewgpuStats: { ...stats, steps: multiStep.steps, multiStepOff } }); + +const multiStep = installMultiStepDecoding(handler.engine, { + steps: DEFAULT_DECODE_STEPS, + onBurst: (b) => { + stats.bursts += 1; + stats.tokens += b.tokens; + stats.encodeMs += b.encodeMs ?? 0; + stats.syncMs += b.syncMs ?? 0; + stats.dispatches += b.dispatches ?? 0; + stats.forwardDispatches += b.forwardDispatches ?? 0; + stats.flushes += b.flushes ?? 0; + postStats(); + }, + // The only message that can ever report this. When the fast path is off there + // are no bursts, so `onBurst` never fires and the decode probe simply stops + // arriving — indistinguishable, from the host's side, from an idle engine. + onFallback: ({ missing }) => { + multiStepOff = missing; + postStats(); + }, +}); + +/** + * Force a full re-prefill instead of reusing the KV cache across rounds. + * + * Multi-round reuse routes attention through `batch_prefill_paged_kv_kernel`, + * which binds 10 storage buffers: q, pages, lse, output and six small i32 + * metadata arrays. Firefox's Metal backend caps `maxStorageBuffersPerShaderStage` + * at 9, so that pipeline fails to build — and an invalid WebGPU pipeline is + * silent, its dispatches becoming no-ops. The symptom is a second turn that + * answers the *previous* question behind a garbage prefix that differs run to + * run, which is uninitialised memory being read. + * + * Resetting the conversation first makes WebLLM's own conversation comparison + * fail, so it re-prefills from scratch through `batch_prefill_ragged_kv_kernel` + * (9 bindings, works). The cost is re-reading the history each turn; prefill is + * one sync per chunk, so it is far cheaper than the garbage it replaces. + * + * Conditional on the limit, not on the browser: a device that allows 10 keeps + * the KV cache and the faster path. + */ +/** Bindings `batch_prefill_paged_kv_kernel` needs; see tools/audit-wasm.mjs. */ +const PAGED_PREFILL_STORAGE_BUFFERS = 10; + +/** Whether this device is too tight to build that pipeline. Probed once. */ +const kvReuseUnsafe = (async () => { + const adapter = await navigator.gpu?.requestAdapter().catch(() => null); + const limit = adapter?.limits?.maxStorageBuffersPerShaderStage ?? 0; + const unsafe = limit < PAGED_PREFILL_STORAGE_BUFFERS; + if (unsafe) { + console.info( + `[everything-webgpu] KV reuse disabled: paged prefill needs ` + + `${PAGED_PREFILL_STORAGE_BUFFERS} storage buffers, this device allows ${limit}`, + ); + } + return unsafe; +})(); + +// Wrapped synchronously at module load: deciding first and wrapping after the +// await would leave a window where an early prefill slips through unguarded. +const basePrefill = handler.engine.prefill.bind(handler.engine); +handler.engine.prefill = async (input, pipeline, chatConfig, genConfig) => { + if (await kvReuseUnsafe) pipeline.resetChat(/* keepStats= */ true); + return basePrefill(input, pipeline, chatConfig, genConfig); +}; + +// The engine is reachable before any model is loaded, so the host can set the +// step count on the very first message and never has to reload to change it. +self.onmessage = (msg) => { + if (msg.data?.kind === WORKER_CONFIGURE) { + if (msg.data.decodeSteps !== undefined) multiStep.setSteps(msg.data.decodeSteps); + // Every retune starts a fresh measurement window, so a sweep's points never + // bleed into each other. + resetStats(); + postStats(); + return; + } + handler.onmessage(msg); +}; diff --git a/src/engine/engine.js b/src/engine/engine.js new file mode 100644 index 0000000000000000000000000000000000000000..ed0662154d91d82a68548a574c8b2eacbfe3def1 --- /dev/null +++ b/src/engine/engine.js @@ -0,0 +1,1336 @@ +/** + * The engine, with no transport attached. + * + * This is what a developer embedding a local model in their own page holds: + * one object, plain async methods, no message vocabulary. `browser.runtime` is + * not referenced anywhere in this file or anything it imports — the WebExtension + * message/port router that used to live here now sits behind + * `src/adapters/webext.js` and calls these same methods. + * + * Everything runs through one EnginePool, which owns priority, cancellation and + * fan-out. Nothing here decides what runs when. + * + * Three things are injected, and they are exactly the three places the host + * environment leaked into the engine: + * + * - `store` a ModelStore over a StorageAdapter, because + * `browser.storage.local` does not exist in a page. + * - `workerUrl` defaults to `new URL("./engine-worker.js", import.meta.url)`, + * which Vite, webpack 5 and esbuild all understand and which + * also resolves correctly on `moz-extension://`. It replaces + * `browser.runtime.getURL` rather than sitting beside it. + * - `loadWebLLM` defaults to a static relative `import()`, so bundlers can see + * it, and a host with its own patched build can override it. + * It stays *dynamic* on purpose: the ~6 MB WebLLM bundle is + * fetched when a model is loaded, not when the page paints. + * + * Weights arrive by any of three routes — prebuilt (WebLLM's HuggingFace list), + * remote (`registerModel` with any base URL you host), or injected + * (`ingestModelFolder`, no network at all). `load()` resolves across all three; + * see `model-store.js`. Pass `prebuilt: false` for a build that must never + * reach the network for a model. + */ +import { ENGINE_STATE, JOB_KIND, PRIORITY, UNLOAD_LEVEL, WORKER_CONFIGURE } from "./constants.js"; +import { chatFacade } from "./chat.js"; +import { environmentFacade } from "./environment.js"; +import { canRun, probeDevice, projectSpeed, rankModels } from "./device.js"; +import { ERROR, EngineError, asEngineError } from "./errors.js"; +import { filesFromDataTransfer, filesFromInput, ingestModelFolder } from "./ingest.js"; +import { ModelStore, SOURCE, groupKeysByScope, isInjected, toAppConfig } from "./model-store.js"; +import { clampSteps } from "./multistep.js"; +import { EnginePool } from "./pool.js"; +import { prefetchModel } from "./prefetch.js"; +import { ask, conversation, ghostText } from "./recipes.js"; +import { SOURCE_KIND, classifySource, isDataTransfer, isFileList, nearMatches } from "./sources.js"; + +/** + * @typedef {object} CompletionRequest + * The OpenAI generation fields WebLLM already speaks, plus the scheduling + * fields that are what this engine adds over calling WebLLM directly. + * @property {Array<{role: string, content: string}>} messages + * @property {string} [modelId] load this model first if it is not the live one + * @property {string} [id] job id; also what `cancel(id)` takes + * @property {number} [temperature] + * @property {number} [max_tokens] + * @property {object} [response_format] + * @property {object} [extra_body] + * @property {string} [task] the unit that owns an engine; a whole batch shares one + * @property {string} [session] a later job with this key supersedes the earlier one + * @property {"interactive"|"normal"|"background"} [priority] + * @property {boolean} [preemptible] may be interrupted by an `interactive` job + */ + +/** + * @typedef {object} CompletionResult + * @property {string} text + * @property {object} [usage] + * @property {"stop"|"length"|"abort"} [finishReason] WebLLM's own values + * @property {true} [cancelled] superseded or explicitly cancelled + * @property {true} [preempted] an `interactive` job took the slot; `text` is partial + */ + +/** + * @typedef {CompletionRequest & {index: number, engineIndex: number, + * startedAt: number, finishedAt: number, error?: string}} BatchItem + */ + +const DEFAULT_WORKER_URL = () => new URL("./engine-worker.js", import.meta.url); +const DEFAULT_LOAD_WEBLLM = () => import("../../vendor/web-llm.js"); + +/** + * Turn "the package is not wired into your build" into a sentence that says so. + * + * `vendor/web-llm.js` is a **build product**, not a checked-in file, so the two + * ways to arrive here are both install-shaped rather than runtime-shaped: a git + * dependency whose `prepare` never ran, or a source checkout where `npm run + * build` was never run. Left alone this surfaced as a bare + * `Cannot find module '.../vendor/web-llm.js'` under the code + * `GENERATION_FAILED` — wrong twice over, since nothing had begun generating + * and the path named is ours, not the caller's. + */ +async function loadBundle(loadWebLLM) { + try { + return await loadWebLLM(); + } catch (err) { + const message = String(err?.message ?? err); + // Only a resolution failure means "not built". A bundle that throws while + // *evaluating* is a real crash and must keep its own stack. + if (!/Cannot find module|Failed to (fetch|resolve)|ERR_MODULE_NOT_FOUND|dynamically imported module/i.test(message)) { + throw err; + } + throw new EngineError( + ERROR.PACKAGE_INCOMPLETE, + "everything-webgpu is installed but its WebLLM bundle (vendor/web-llm.js) is missing. " + + "That file is generated, not checked in — run `npm run build` in the package, " + + "or reinstall so its `prepare` script runs.", + { cause: "vendor-bundle-missing", underlying: message }, + ); + } +} + +export class ScheduledEngine { + #store; + #workerUrl; + #loadWebLLM; + #prebuilt; + #chat = null; + #environment = null; + #probe = null; + /** This machine's achieved decode bandwidth, learned from the first generation. */ + #decodeBytesPerSecond = 0; + /** modelId -> weight bytes, for projections. */ + #modelBytes = new Map(); + /** + * Resident models: modelId -> EnginePool. More than one may be up at once — + * a text model beside a vision model, say — which is why this is a map and + * not a field. Each entry holds a full copy of its weights, so residency is + * budget-gated in `load()`. + * @type {Map} + */ + #pools = new Map(); + /** Which resident model an unaddressed request goes to. */ + #current = null; + /** In-flight loads, per model, so two models can come up concurrently. */ + #loading = new Map(); + /** + * Pools that are still loading. `load()` holds the pool in a local until it + * is ready, which left an in-flight download unreachable — so aborting one + * had nothing to tear down. See `load({ signal })`. + * @type {Map} + */ + #loadingPools = new Map(); + #listeners = new Set(); + + #state = { + status: ENGINE_STATE.IDLE, + modelId: null, + progress: null, + error: null, + pool: { size: 0, busy: 0, queued: 0 }, + /** Model ids with a live pool. `modelId` is whichever of them is current. */ + resident: [], + /** Latest decode probe from an engine worker; see multistep.js. */ + decode: null, + }; + + /** + * @param {object} opts + * @param {ModelStore | import("./model-store.js").StorageAdapter} opts.store + * a ModelStore, or a bare StorageAdapter to wrap in one + * @param {string | URL} [opts.workerUrl] + * @param {() => Promise} [opts.loadWebLLM] + * @param {boolean} [opts.prebuilt] expose WebLLM's 163 HuggingFace-hosted + * models, downloaded on first load. Default true. Set false for an + * offline-only build: `load()` then resolves registered models and nothing + * else, and an unknown id fails before the WebLLM bundle is even fetched. + */ + constructor({ store, workerUrl, loadWebLLM, prebuilt = true } = {}) { + if (!store) { + throw new EngineError(ERROR.BAD_REQUEST, "ScheduledEngine needs a `store` (ModelStore or StorageAdapter)."); + } + this.#store = store instanceof ModelStore ? store : new ModelStore(store); + this.#workerUrl = workerUrl ?? DEFAULT_WORKER_URL(); + // Wrapped once here rather than at each of the seven `#loadWebLLM()` call + // sites: a missing bundle is the same failure whichever verb reached it + // first, and a site added later gets the good error for free. + const load = loadWebLLM ?? DEFAULT_LOAD_WEBLLM; + this.#loadWebLLM = () => loadBundle(load); + this.#prebuilt = prebuilt; + } + + /** The ModelStore, so a host can drive the registry without a second handle. */ + get store() { + return this.#store; + } + + /** + * `chat.completions.create()`, the WebLLM/OpenAI shape. See `chat.js`. + * + * Built once and cached: callers hold on to `engine.chat.completions` the way + * they did with WebLLM, and a fresh object each access would break that. + */ + get chat() { + this.#chat ??= chatFacade(this); + return this.#chat; + } + + /** + * `environment()` — the read-only report, with `environment.measure()` on it. + * + * Cached like `chat` so a caller can hold on to it. Writes are `configure()`; + * see `environment.js` for why those are separate verbs. + */ + get environment() { + this.#environment ??= environmentFacade(this); + return this.#environment; + } + + get state() { + return { ...this.#state }; + } + + get hasWebGPU() { + return Boolean(globalThis.navigator?.gpu); + } + + /** + * @param {(state: object) => void} listener called immediately, then on change + * @returns {() => void} unsubscribe + */ + subscribe(listener) { + this.#listeners.add(listener); + listener(this.state); + return () => this.#listeners.delete(listener); + } + + /** The current model's pool, or null. */ + get #pool() { + return this.#current ? (this.#pools.get(this.#current) ?? null) : null; + } + + /** Model ids with a live pool right now. */ + get resident() { + return [...this.#pools.keys()]; + } + + /** + * Choose which resident model unaddressed requests go to. + * + * Distinct from `load()` on purpose: this is free and instant, because the + * weights are already up. `load()` is what costs. + */ + use(modelId) { + if (!this.#pools.has(modelId)) { + throw new EngineError( + ERROR.UNKNOWN_MODEL, + `"${modelId}" is not resident. Resident: ${this.resident.join(", ") || "none"}. Call load() first.`, + { modelId, resident: this.resident }, + ); + } + this.#current = modelId; + this.#syncState(); + return this.state; + } + + /** Registered models only — cheap, no bundle load. */ + listModels() { + return this.#store.list(); + } + + /** + * Everything `load()` would accept, normalised: registered models first, then + * WebLLM's prebuilt list. + * + * Costs a WebLLM bundle fetch when `prebuilt` is on, because the list lives + * inside it. `listModels()` is the cheap call if you only care about what this + * app registered. + * + * @returns {Promise>} + */ + async listAvailableModels() { + const registered = await this.#store.list(); + const own = registered.map((r) => ({ + modelId: r.model_id, + source: r.source ?? SOURCE.REMOTE, + model: r.model, + contextWindow: r.overrides?.context_window_size, + vramRequiredMB: r.vram_required_MB, + sizeBytes: r.sizeBytes, + })); + if (!this.#prebuilt) return own; + + const { prebuiltAppConfig, functionCallingModelIds } = await this.#loadWebLLM(); + const toolCalling = new Set(functionCallingModelIds ?? []); + const owned = new Set(own.map((m) => m.modelId)); + const rest = prebuiltAppConfig.model_list + .filter((e) => !owned.has(e.model_id)) + .map((e) => ({ + modelId: e.model_id, + source: SOURCE.PREBUILT, + model: e.model, + contextWindow: e.overrides?.context_window_size, + vramRequiredMB: e.vram_required_MB, + // WebLLM ships the list; it is not derivable from the id. + toolCalling: toolCalling.has(e.model_id), + })); + return [...own, ...rest]; + } + + /** + * What this machine will admit to: WebGPU, adapter, `shader-f16`, the five + * limits that matter, storage quota. Cached — hardware does not change + * mid-session, and `requestAdapter()` is not free. + * @returns {Promise} + */ + probe() { + this.#probe ??= probeDevice(); + return this.#probe; + } + + /** + * Whether a model will run here, before anything is downloaded. + * @param {string} modelId + * @returns {Promise<{ok: boolean, blockers: Array, warnings: Array}>} + */ + async canRun(modelId) { + const [probe, available] = await Promise.all([this.probe(), this.listAvailableModels()]); + const found = available.find((m) => m.modelId === modelId); + if (!found) { + throw new EngineError(ERROR.UNKNOWN_MODEL, `Model "${modelId}" is not one this engine can load.`, { + modelId, + }); + } + return canRun( + { model_id: found.modelId, vram_required_MB: found.vramRequiredMB, sizeBytes: found.sizeBytes }, + probe, + ); + } + + /** + * Which models this device should actually be asked to run, best first. + * + * The prebuilt list spans 239 MB to 31 GB; this is the answer to the first + * question a developer has and the one they have least basis to answer. + * + * @param {{maxVramMB?: number, needsVision?: boolean, needsToolCalling?: boolean, + * prefer?: "quality" | "speed"}} [opts] + */ + async recommendModels({ needsToolCalling = false, ...opts } = {}) { + const [probe, appConfig] = await Promise.all([this.probe(), this.#appConfig()]); + let list = appConfig.model_list; + if (needsToolCalling) { + const { functionCallingModelIds } = await this.#loadWebLLM(); + const ids = new Set(functionCallingModelIds ?? []); + list = list.filter((m) => ids.has(m.model_id)); + } + return rankModels(list, { probe, ...opts }); + } + + /** + * Is this model's data on disk, so a load would need no network? + * + * Routes by who knows the keys. We wrote an injected model's artifacts and + * hold the manifest, so `verify()` answers exactly — including a `"partial"` + * verdict WebLLM cannot give. Everything else was fetched by WebLLM, which + * derives the keys as its loader did, so `hasModelInCache` is the answer. + * + * @returns {Promise<"cached" | "partial" | "absent">} + */ + async cacheState(modelId) { + const record = await this.#store.get(modelId); + if (isInjected(record)) { + const { ok, missing } = await this.#store.verify(record); + if (ok) return "cached"; + const total = Object.values(groupKeysByScope(record)).flat().length; + return missing.length >= total ? "absent" : "partial"; + } + const { hasModelInCache } = await this.#loadWebLLM(); + return (await hasModelInCache(modelId, await this.#appConfig())) ? "cached" : "absent"; + } + + /** + * Download a model into the cache **without building an engine**. + * + * For warming during onboarding: the bytes land while the user is still + * reading, and the later `load()` is a cache read. WebLLM cannot express this + * — `reload()` instantiates the wasm and needs a GPU before it fetches a + * single shard — so this is ours. See `prefetch.js` for the URL-derivation + * risk and the oracle that closes it. + * + * Needs no WebGPU at all, which is the other half of the point: an app can + * warm the cache on a machine it has not yet decided can run the model. + * + * @param {string} modelId + * @param {{signal?: AbortSignal, onProgress?: Function}} [opts] + */ + async prefetch(modelId, { signal, onProgress } = {}) { + // Before anything, including the already-cached shortcut: a caller who + // aborted wants an abort, not a success they have to inspect to distrust. + if (signal?.aborted) { + throw new EngineError(ERROR.ABORTED, `Prefetch of "${modelId}" was aborted before it began.`, { + modelId, + }); + } + const record = await this.#store.get(modelId); + + // An injected model's bytes were written before it was ever registered; + // there is no URL to fetch from and nothing to do. + if (isInjected(record)) { + const { ok } = await this.#store.verify(record); + if (ok) return { modelId, files: 0, bytes: 0, alreadyCached: true }; + throw new EngineError( + ERROR.CACHE_INCOMPLETE, + `"${modelId}" was injected from a folder, so it cannot be re-fetched. Re-register the folder.`, + { modelId }, + ); + } + + const appConfig = await this.#appConfig(); + const entry = appConfig.model_list.find((m) => m.model_id === modelId); + if (!entry) { + const near = nearMatches(modelId, appConfig.model_list.map((m) => m.model_id)); + throw new EngineError( + ERROR.UNKNOWN_MODEL, + `Cannot prefetch "${modelId}": it is neither registered nor prebuilt. ` + + (near.length ? `Did you mean ${near.map((id) => `"${id}"`).join(", ")}? ` : ""), + { modelId, ...(near.length ? { near } : {}) }, + ); + } + + const { hasModelInCache } = await this.#loadWebLLM(); + if (await hasModelInCache(modelId, appConfig).catch(() => false)) { + return { modelId, files: 0, bytes: 0, alreadyCached: true }; + } + + const result = await prefetchModel({ modelId, record: entry, signal, onProgress }); + + // The oracle. `hasModelInCache` derives its keys through the very function + // `prefetch.js` mirrors, so this is the one check that can tell a correct + // prefetch from one that wrote a cache the loader will never read. Without + // it, a wrong key looks exactly like success and costs the user a second + // download of the whole model. + if (!(await hasModelInCache(modelId, appConfig).catch(() => false))) { + throw new EngineError( + ERROR.GENERATION_FAILED, + `Prefetch of "${modelId}" wrote ${result.files} artifacts, but WebLLM still reports the ` + + "model as uncached — the derived cache keys do not match the ones its loader looks for. " + + "Treat the cache as cold; load() will re-download. This is what a change to WebLLM's " + + "URL scheme looks like from here.", + { modelId, files: result.files }, + ); + } + return result; + } + + /** + * Free a model's bytes and **keep the registry entry**, so it stays a model + * this engine knows how to get again — the distinction from + * `store.remove()`, which forgets the URL a remote model would need. + * + * Delegates for remote and prebuilt models: `deleteModelAllInfoInCache` is + * WebLLM's, covers tensors + wasm + config, and is maintained upstream. + */ + async evict(modelId) { + if (this.#pools.has(modelId)) await this.unload(modelId); + return this.#evictBytes(modelId); + } + + /** + * The byte-freeing half of `evict()`, with no pool handling. + * + * Split out so `unload(id, "cache")` can reach it without going back through + * `evict()` → `unload()`, which would re-enter this class for a pool that has + * just been torn down. + */ + async #evictBytes(modelId) { + const record = await this.#store.get(modelId); + if (isInjected(record)) return this.#store.evictInjected(modelId); + + const { deleteModelAllInfoInCache } = await this.#loadWebLLM(); + await deleteModelAllInfoInCache(modelId, await this.#appConfig()); + return { freedKeys: null }; + } + + /** The merged model list WebLLM's cache helpers key off. */ + async #appConfig() { + const registered = await this.#store.list(); + const prebuilt = this.#prebuilt ? (await this.#loadWebLLM()).prebuiltAppConfig : null; + return toAppConfig(registered, prebuilt); + } + + /** + * Forget a model entirely: free its bytes **and** drop the registry entry. + * + * `evict()` first, because that is what knows how to reach the bytes for each + * source — and it has to happen before the record is deleted, since for a + * remote model the record holds the only URL those bytes can be derived from. + * Deleting the entry first would strand them in Cache Storage permanently. + */ + async remove(modelId) { + const freed = await this.evict(modelId); + await this.#store.remove(modelId); + return freed; + } + + /** + * Projected decode throughput for a model, in tokens per second. + * + * `basis: "measured"` once anything has actually decoded on this machine — + * the engine then knows its own achieved bandwidth and every projection is + * device-specific. Before that, `basis: "extrapolated"` from a reference + * machine, which is a starting point and says so. + * + * Decode is memory-bandwidth-bound, so this is close to the whole story: + * time per token scales with weight bytes and little else. + * + * @param {string} [modelId] defaults to the current model + */ + async estimateSpeed(modelId = this.#current) { + let bytes = modelId ? this.#modelBytes.get(modelId) : 0; + if (!bytes && modelId) { + const found = (await this.listAvailableModels()).find((m) => m.modelId === modelId); + bytes = found?.sizeBytes ?? (found?.vramRequiredMB ?? 0) * 1024 * 1024; + } + return { modelId, ...projectSpeed(bytes, this.#decodeBytesPerSecond) }; + } + + /** + * What is actually switched on right now, as opposed to what the device could + * support. + * + * The distinction matters for KV reuse in particular: `probe().kvReuse` is a + * device capability, but the decision is taken inside the engine worker, + * which is the authority. A caller debugging "why is my second turn slow" + * needs the decision, not the capability. + */ + async features() { + const probe = await this.probe(); + const settings = await this.#store.getSettings(); + const pool = this.#pool?.status(); + // The configured K, which is not the same question as whether bursts run. + // `decodeSteps` stays the knob's value even when the fast path is off: it is + // what `configure()` would be changing, and reporting 1 here would make + // `environment()` advise `configure({ decodeSteps: 15 })` for a fault that + // setting cannot fix. + const decodeSteps = this.#state.decode?.steps ?? settings.decodeSteps; + // Non-empty means the live pipeline failed the contract and decoding was + // routed to stock single-step. The worker keeps posting the *configured* K + // either way, so `steps > 1` cannot answer this — asking it alone reported + // a disabled fast path as healthy, which is the one case this call exists for. + const multiStepOff = this.#state.decode?.multiStepOff; + return { + kvReuse: Boolean(probe.kvReuse), + shaderF16: Boolean(probe.features?.shaderF16), + decodeSteps, + multiStepDecoding: decodeSteps > 1 && !multiStepOff?.length, + /** The missing tvmjs internals, or null when the fast path is intact. */ + multiStepOff: multiStepOff?.length ? multiStepOff : null, + engines: pool?.size ?? 0, + maxEngines: pool?.maxSize ?? settings.engineCount, + resident: this.resident, + // dispatches per flush: >1 means consecutive kernel launches are sharing a + // compute pass, i.e. the build-time batching patch is in effect. Observed, + // not asserted — a build with NO_PASS_MERGE=1 reports ~1. + computePassBatching: this.#state.decode?.flushes + ? this.#state.decode.dispatches / this.#state.decode.flushes + : null, + decode: this.#state.decode ?? null, + }; + } + + /** + * Register a model. Two shapes, one call, and the difference is only where + * the bytes come from: + * + * ```js + * // fetched from a base URL you host — an HF repo, a CDN, your own origin + * await engine.registerModel({ + * modelId: "my-model", + * model: "/models/my-model/", + * modelLib: "/models/my-model/my-model-webgpu.wasm", + * }); + * + * // read off disk. No network connection at any point, ever. + * await engine.registerModel({ modelId: "my-model", files: entries }); + * ``` + * + * Both end up as one `model_list` entry that WebLLM's own loader resolves the + * same way — the local one only differs in that its base URL is minted on + * `.invalid` and its cache is populated before the loader ever looks. + * + * That origin is the *mechanism* of the offline guarantee, not a marker of + * it: `.invalid` is reserved by RFC 6761 and can never resolve, so there is + * no code path — no bug, no eviction, no future refactor — by which a local + * model reaches the network. It fails with a DNS error instead. + * + * `files` is `{ path, file }[]`; `filesFromDataTransfer` and + * `filesFromInput` build it from a drop event or a directory picker. + */ + async registerModel(spec) { + if (spec?.files) { + if (spec.model || spec.modelLib) { + throw new EngineError( + ERROR.BAD_REQUEST, + "registerModel takes either `files` (local, never fetched) or `model`/`modelLib` (a base URL to fetch), not both.", + ); + } + return ingestModelFolder(spec.files, { + store: this.#store, + modelId: spec.modelId, + modelType: spec.modelType, + onProgress: spec.onProgress, + }); + } + return this.#store.registerModel(spec); + } + + // ---------------------------------------------------------------- engine --- + + #setState(patch) { + Object.assign(this.#state, patch); + const snapshot = this.state; + for (const listener of [...this.#listeners]) { + try { + listener(snapshot); + } catch { + /* a subscriber that throws must not stall the engine */ + } + } + } + + #assertWebGPU() { + if (!this.hasWebGPU) { + throw new EngineError( + ERROR.NO_WEBGPU, + "WebGPU is unavailable in this context. On macOS Firefox, set dom.webgpu.enabled=true " + + "(and gfx.webgpu.ignore-blocklist=true if your GPU is blocklisted) in about:config, then restart Firefox.", + ); + } + } + + /** + * Bring a model up, whatever form you have it in. + * + * One entry point for all three routes, because from a caller's side "load a + * model" is one intention and having to know which of `load`, + * `registerModel` and `ingestModelFolder` to reach for is a decision the + * library can make for them: + * + * ```js + * load("Llama-3.2-1B-Instruct-q4f16_1-MLC") // prebuilt or registered id + * load("https://huggingface.co/mlc-ai/Foo", { modelLib }) // a URL you host + * load({ model, modelLib }) // the same, explicit + * load({ files }) | load(fileList) | load(dataTransfer) // a folder, no network + * ``` + * + * `registerModel` and `ingestModelFolder` remain, unchanged, as the low-level + * primitives — this composes them rather than replacing them. + * + * **A URL always needs `modelLib`.** It is not guessed; see `sources.js` for + * the measurement behind that. **`defer: true`** registers the source and + * stops there, returning the record instead of the state — the manager's + * drop-now-load-later flow. + * + * Additive residency: a model already resident stays resident, so switching + * back to it costs nothing. That is only safe while the weights fit, so + * `keepResident: false` (the default) unloads whatever else is up first — + * the old single-model behaviour, and the safe one on a 16 GB machine. + * Pass `keepResident: true` to hold both, having checked the budget yourself + * with `canRun()`. + * + * @param {string | object} src an id, a URL, `{model, modelLib}`, or a folder + * @param {{keepResident?: boolean, signal?: AbortSignal, defer?: boolean, + * id?: string, modelLib?: string, modelType?: string, contextWindow?: number, + * vramRequiredMB?: number, onProgress?: Function}} [opts] + * @returns {Promise} the engine state, or the registry record when `defer` + */ + async load(src, opts = {}) { + const source = classifySource(src, opts); + + if (source.kind === SOURCE_KIND.ID) { + if (opts.defer) { + throw new EngineError( + ERROR.BAD_REQUEST, + `\`defer\` registers a source without loading it, but "${source.modelId}" is an id — ` + + "there is nothing to register. Drop `defer`, or pass a URL or a folder.", + { modelId: source.modelId }, + ); + } + return this.#loadById(source.modelId, opts); + } + + const record = await this.#register(source, opts); + if (opts.defer) return record; + return this.#loadById(record.model_id, opts); + } + + /** Turns a classified non-id source into a registry record. */ + async #register(source, opts) { + if (source.kind === SOURCE_KIND.FILES) { + return ingestModelFolder(await toEntries(source.files), { + store: this.#store, + modelId: source.modelId, + modelType: opts.modelType, + onProgress: opts.onProgress, + }); + } + return this.#store.registerModel({ + modelId: source.modelId, + model: source.model, + modelLib: source.modelLib, + modelType: opts.modelType, + contextWindow: opts.contextWindow, + vramRequiredMB: opts.vramRequiredMB, + }); + } + + /** + * Bring a registered or prebuilt id up and make it current. + * + * Cancellation is WebLLM's, not ours: `unload()` aborts the `reloadController` + * whose signal it threads through every artifact fetch. And resume is free — + * WebLLM caches each artifact as it arrives and skips what is already present, + * so a re-`load()` picks up where the abort left off. Call `evict()` to + * discard a partial download instead. + * + * @param {string} modelId + * @param {{keepResident?: boolean, signal?: AbortSignal}} [opts] + */ + async #loadById(modelId, { keepResident = false, signal } = {}) { + if (signal?.aborted) { + throw new EngineError(ERROR.ABORTED, `Load of "${modelId}" was aborted before it began.`, { + modelId, + }); + } + if (this.#pools.has(modelId)) return this.use(modelId); + + const inFlight = this.#loading.get(modelId); + if (inFlight) { + await inFlight.catch(() => {}); + return this.#pools.has(modelId) ? this.use(modelId) : this.state; + } + + const loading = (async () => { + this.#assertWebGPU(); + + const models = await this.#store.list(); + const registered = models.find((m) => m.model_id === modelId); + + // Only an injected model can be *unrecoverably* broken by eviction, and + // checking costs nothing, so it gates the load — before the WebLLM bundle + // is fetched, so the common offline failure stays fast. + if (isInjected(registered)) { + const { ok, missing } = await this.#store.verify(registered); + if (!ok) { + throw new EngineError( + ERROR.CACHE_INCOMPLETE, + `Cache for "${modelId}" is incomplete (${missing.length} artifact(s) evicted, e.g. ${missing[0].split("/").pop()}). Re-register the model folder.`, + { modelId, missing }, + ); + } + } + + if (!registered && !this.#prebuilt) { + const near = nearMatches(modelId, models.map((m) => m.model_id)); + throw new EngineError( + ERROR.UNKNOWN_MODEL, + `Model "${modelId}" is not registered, and prebuilt models are disabled. ` + + (near.length ? `Did you mean ${near.map((id) => `"${id}"`).join(", ")}? ` : "") + + `Call load(url, { modelLib }) or load({ files }) first.`, + { modelId, prebuilt: false, ...(near.length ? { near } : {}) }, + ); + } + + const { engineCount, decodeSteps } = await this.#store.getSettings(); + this.#setState({ + status: ENGINE_STATE.LOADING, + modelId, + error: null, + progress: { text: "Starting", progress: 0 }, + }); + + // Each resident model is a full copy of its weights. Nothing reports free + // VRAM to a page (AI.md, "The pool grows, it is not sized"), so the + // default is to make room rather than to gamble on it fitting. + if (!keepResident) { + for (const id of [...this.#pools.keys()]) await this.#unloadOne(id); + } + + const { CreateWebWorkerMLCEngine, prebuiltAppConfig } = await this.#loadWebLLM(); + const appConfig = toAppConfig(models, this.#prebuilt ? prebuiltAppConfig : null); + + if (!appConfig.model_list.some((m) => m.model_id === modelId)) { + // A typo'd id is the single most likely way to arrive here, and the + // fix is almost always visible in the list we are already holding. + const near = nearMatches(modelId, appConfig.model_list.map((m) => m.model_id)); + throw new EngineError( + ERROR.UNKNOWN_MODEL, + `Model "${modelId}" is neither registered nor in WebLLM's prebuilt list. ` + + (near.length + ? `Did you mean ${near.map((id) => `"${id}"`).join(", ")}? ` + : "") + + `Use listAvailableModels() to see what this engine can load.`, + { modelId, ...(near.length ? { near } : {}) }, + ); + } + + const pool = new EnginePool({ + size: engineCount, + createEngine: async (_index, onProgress) => { + const worker = new Worker(this.#workerUrl, { type: "module" }); + // Listener, not `onmessage`: WebLLM claims `onmessage` for its own RPC. + worker.addEventListener("message", (event) => { + if (event.data?.ewgpuStats) this.#setState({ decode: event.data.ewgpuStats }); + }); + // A worker whose script 404s does not throw from `new Worker()` — it + // fires one `error` event and is then simply silent, so WebLLM's + // handshake below never resolves and the load hangs until the caller + // gives up. That is the exact shape of the Vite dep-optimizer bug this + // names: esbuild copies `new URL("./engine-worker.js", + // import.meta.url)` into `.vite/deps/` verbatim, where the sibling + // file does not exist. Racing the handshake against this turns a hang + // into a sentence. + const workerFailed = new Promise((_, reject) => { + worker.addEventListener("error", (event) => { + reject( + new EngineError( + ERROR.PACKAGE_INCOMPLETE, + `The decode worker failed to load from ${this.#workerUrl}. ` + + "If you are on Vite, its dependency pre-bundler rewrote the worker URL to a " + + "path that does not exist — add `optimizeDeps: { exclude: [\"everything-webgpu\"] }` " + + "to vite.config.js, or pass `workerUrl` yourself.", + { cause: "worker-unreachable", workerUrl: String(this.#workerUrl), underlying: event.message }, + ), + ); + }); + }); + // Sent before WebLLM's own handshake so the first token already decodes + // multi-step; worker message order guarantees it arrives first. + worker.postMessage({ kind: WORKER_CONFIGURE, decodeSteps }); + const engine = await Promise.race([ + CreateWebWorkerMLCEngine(worker, modelId, { + appConfig, + initProgressCallback: onProgress, + }), + workerFailed, + ]); + // The worker owns the decode loop, so runtime knobs go straight to it + // rather than through WebLLM's request path. + engine.configure = (patch) => worker.postMessage({ kind: WORKER_CONFIGURE, ...patch }); + // Tear the realm down with the engine, not just the model. + const unloadEngine = engine.unload.bind(engine); + engine.unload = async () => { + await unloadEngine().catch(() => {}); + worker.terminate(); + }; + return engine; + }, + onStateChange: () => { + if (this.#current === modelId) this.#syncState(); + }, + }); + + // Reachable from here on, so an abort has something to tear down. The + // listener bumps the pool's generation, which is what stops an engine + // that finishes building *after* the abort from installing itself. + this.#loadingPools.set(modelId, pool); + if (signal) signal.addEventListener("abort", () => void this.#unloadOne(modelId), { once: true }); + + const abortedNow = async () => { + await this.#unloadOne(modelId); + throw new EngineError(ERROR.ABORTED, `Load of "${modelId}" was aborted.`, { modelId }); + }; + + // Checked on both sides of the load. Before: an abort that landed while + // the registry was being read must not start a download at all. After: an + // abort during the download has already torn the pool down, and this is + // what turns that into a rejection rather than a silent no-op. + if (signal?.aborted) return abortedNow(); + await pool.load((progress) => this.#setState({ progress })); + if (signal?.aborted) return abortedNow(); + + const entry = appConfig.model_list.find((m) => m.model_id === modelId); + this.#modelBytes.set( + modelId, + registered?.sizeBytes ?? (entry?.vram_required_MB ?? 0) * 1024 * 1024, + ); + this.#pools.set(modelId, pool); + this.#current = modelId; + this.#setState({ status: ENGINE_STATE.READY, progress: null, error: null }); + this.#syncState(); + return this.state; + })(); + + this.#loading.set(modelId, loading); + try { + return await loading; + } catch (err) { + const engineError = asEngineError(err); + this.#setState({ + status: this.#pools.size ? ENGINE_STATE.READY : ENGINE_STATE.ERROR, + progress: null, + error: engineError.message, + errorCode: engineError.code, + }); + // A failed load must not leave `modelId` pointing at the model that did + // not come up — it was set optimistically when LOADING began. + this.#syncState(); + throw engineError; + } finally { + this.#loading.delete(modelId); + this.#loadingPools.delete(modelId); + } + } + + /** + * Let a model go, at one of two depths. + * + * ```js + * unload() // the current model's VRAM; cached bytes stay + * unload(id) // that model's VRAM + * unload(id, "cache") // and delete its cached bytes, keeping the registry entry + * ``` + * + * At `"vram"` the bytes stay on disk, so loading it again costs no network — + * that is what makes switching back cheap, and the difference between this + * and `remove()`. + * + * **A bare `unload()` frees only the current model**, not every resident one. + * `unloadAll()` is the explicit form for that: freeing everything is the more + * destructive of the two readings and should have to be asked for by name. + * + * @param {string} [modelId] defaults to the current model. Omit both this and + * any resident model to no-op. + * @param {"vram"|"cache"} [level] + */ + async unload(modelId = this.#current, level = UNLOAD_LEVEL.VRAM) { + if (!Object.values(UNLOAD_LEVEL).includes(level)) { + throw new EngineError( + ERROR.BAD_REQUEST, + `unload() level must be ${Object.values(UNLOAD_LEVEL).map((l) => `"${l}"`).join(" or ")}, ` + + `not "${level}". To forget the model entirely, use remove().`, + { level }, + ); + } + if (modelId) { + await this.#unloadOne(modelId); + if (level === UNLOAD_LEVEL.CACHE) await this.#evictBytes(modelId); + } + this.#setState({ + status: this.#pools.size ? ENGINE_STATE.READY : ENGINE_STATE.IDLE, + progress: null, + error: null, + }); + this.#syncState(); + return this.state; + } + + /** Unload every resident model. */ + async unloadAll() { + for (const id of [...this.#pools.keys()]) await this.#unloadOne(id); + this.#setState({ status: ENGINE_STATE.IDLE, progress: null, error: null }); + this.#syncState(); + return this.state; + } + + async #unloadOne(modelId) { + const pool = this.#pools.get(modelId) ?? this.#loadingPools.get(modelId); + if (!pool) return; + this.#pools.delete(modelId); + this.#loadingPools.delete(modelId); + if (this.#current === modelId) this.#current = this.#pools.keys().next().value ?? null; + await pool.unload(); + } + + /** + * Learn this machine's achieved decode bandwidth from a finished generation. + * + * WebLLM already measures throughput and puts it on every response as + * `usage.extra.decode_tokens_per_s` — the pool asks for usage and was + * discarding it. Tokens per second times weight bytes is bytes per second, + * which is the one number a projection for *any other* model needs. + * + * The worker's own decode probe is not this and does not replace it: it + * splits a burst into CPU-encode and GPU-sync, which is what the multi-step + * and compute-pass work is measured against. Nothing upstream reports that. + */ + #calibrate(usage, modelId) { + const tps = usage?.extra?.decode_tokens_per_s; + const bytes = this.#modelBytes.get(modelId); + if (!tps || !bytes) return; + this.#decodeBytesPerSecond = tps * bytes; + } + + /** Recomputes the parts of `state` that are views onto the current pool. */ + #syncState() { + this.#setState({ + modelId: this.#current, + resident: this.resident, + pool: this.#pool?.status() ?? { size: 0, busy: 0, queued: 0 }, + }); + } + + /** Loads on demand so callers can just ask for a completion. */ + async #ensurePool(modelId) { + // A request naming a resident model goes straight to it, without disturbing + // which model is current — routing is not switching. + if (modelId && this.#pools.has(modelId)) return this.#pools.get(modelId); + if (modelId) await this.load(modelId); + if (!this.#pool) { + const fallback = modelId ?? this.#current ?? (await this.#store.list())[0]?.model_id; + if (!fallback) { + throw new EngineError( + ERROR.NO_MODEL, + "No model is registered yet. Call registerModel(), or load() a prebuilt id.", + ); + } + await this.load(fallback); + } + return modelId ? (this.#pools.get(modelId) ?? this.#pool) : this.#pool; + } + + async #buildParams(payload) { + const settings = await this.#store.getSettings(); + const messages = Array.isArray(payload.messages) ? [...payload.messages] : []; + if (messages.length === 0) { + throw new EngineError(ERROR.BAD_REQUEST, "`messages` must be a non-empty array."); + } + if (settings.systemPrompt && !messages.some((m) => m.role === "system")) { + messages.unshift({ role: "system", content: settings.systemPrompt }); + } + return { + messages, + temperature: payload.temperature ?? settings.temperature, + max_tokens: payload.max_tokens ?? settings.maxTokens, + ...(payload.response_format ? { response_format: payload.response_format } : {}), + ...(payload.extra_body ? { extra_body: payload.extra_body } : {}), + }; + } + + /** + * One completion. + * + * Named `complete` rather than `chat` so `engine.chat.completions.create()` + * — the WebLLM-shaped facade, Phase 2 — can take that name without a rename. + * + * @param {CompletionRequest} payload + * @param {(delta: string) => void} [onChunk] called per streamed text delta + * @returns {Promise} + */ + async complete(payload, onChunk) { + // Adapts the raw chunk stream to the documented string callback. The guard + // matters: raw chunks include a role-only first chunk and a finish-only + // last chunk, which callers of this signature have never seen. + return this.completeRaw( + payload, + onChunk && + ((chunk) => { + const delta = chunk?.choices?.[0]?.delta?.content; + if (delta) onChunk(delta); + }), + ); + } + + /** + * `complete()`, but the callback receives WebLLM's chunk verbatim. + * + * Exists so the `chat.completions.create()` facade can pass chunks straight + * through instead of rebuilding an envelope — which is what dropped + * `tool_calls`, flattened `logprobs` and restamped `created`. + * + * @param {CompletionRequest} payload + * @param {(chunk: object) => void} [onRawChunk] + * @returns {Promise}>} + */ + async completeRaw(payload, onRawChunk) { + const pool = await this.#ensurePool(payload.modelId); + const result = unwrap( + await pool.submit({ + ...scheduling(payload), + id: payload.id, + params: await this.#buildParams(payload), + onChunk: onRawChunk, + }), + ); + this.#calibrate(result.usage, this.#current); + return { + text: result.text, + usage: result.usage, + finishReason: result.finishReason, + ...(result.toolCalls ? { toolCalls: result.toolCalls } : {}), + ...(result.cancelled ? { cancelled: true } : {}), + ...(result.preempted ? { preempted: true } : {}), + }; + } + + // ------------------------------------------- the three shapes, as verbs --- + // + // `complete()` expresses all three. These exist because the scheduling is the + // part that is easy to get wrong and invisible when you do — see recipes.js. + + /** + * One question, one answer, nothing kept. + * + * ```js + * const answer = await engine.ask("Summarise this in one line:\n" + doc); + * ``` + * + * @param {string | Array} input + * @param {object} [opts] anything `complete()` takes, plus `onDelta` to stream + * @returns {Promise} + */ + ask(input, opts) { + return ask(this, input, opts); + } + + /** + * A multi-turn conversation that keeps its own history. + * + * ```js + * const chat = engine.conversation({ system: "You are terse." }); + * await chat.say("hello"); + * await chat.say("and again?"); // remembers + * ``` + * + * @param {object} [opts] `system`, `keep`, plus `complete()` defaults + */ + conversation(opts) { + return conversation(this, opts); + } + + /** + * Ghost text, with the debounce/supersede/drop-if-stale discipline built in + * and the prompt left to you. + * + * ```js + * const ghost = engine.ghostText({ prompt: (before) => `Continue:\n${before}` }); + * editor.on("input", async () => { + * const hint = await ghost.suggest(editor.textBefore()); + * if (hint !== null) render(hint); // null means a newer keystroke won + * }); + * editor.on("blur", () => ghost.cancel()); + * ``` + * + * @param {object} opts must include `prompt` + */ + ghostText(opts) { + return ghostText(this, opts); + } + + /** + * Embed text into vectors, through the same scheduler as everything else. + * + * ```js + * const [vector] = await engine.embed("a sentence", { modelId: EMBED_MODEL }); + * const vectors = await engine.embed(["one", "two"], { modelId: EMBED_MODEL }); + * ``` + * + * **Needs an embedding model**, not a chat model — `snowflake-arctic-embed-*` + * in WebLLM's prebuilt list, from 239 MB. They are separate models, so this + * usually names `modelId` explicitly and holds it resident alongside a chat + * model with `load(id, { keepResident: true })`. + * + * Returns bare vectors because that is what a caller does arithmetic on; the + * OpenAI envelope is available as `embedRaw()` for anyone porting code that + * expects `data[].embedding`. + * + * **A running embedding cannot be interrupted.** Cancellation and preemption + * work by making a decode loop break out; one forward pass has no loop, so a + * `cancel()` that lands after the job starts marks it cancelled but does not + * stop it. Queued embeddings supersede and cancel normally. This is tolerable + * because an embedding is milliseconds where a completion is seconds — but it + * is a weaker guarantee than `complete()` gives, so it is stated rather than + * discovered. + * + * @param {string | string[]} input + * @param {{modelId?: string, task?: string, session?: string, + * priority?: string, preemptible?: boolean, id?: string}} [opts] + * @returns {Promise} one vector per input, in order + */ + async embed(input, opts = {}) { + const { data } = await this.embedRaw(input, opts); + return data.map((d) => d.embedding); + } + + /** `embed()`, returning WebLLM's OpenAI-shaped envelope untouched. */ + async embedRaw(input, opts = {}) { + const texts = Array.isArray(input) ? input : [input]; + if (texts.length === 0 || texts.some((t) => typeof t !== "string")) { + throw new EngineError( + ERROR.BAD_REQUEST, + "embed() takes a string or a non-empty array of strings.", + { received: Array.isArray(input) ? `array of ${input.length}` : typeof input }, + ); + } + + const pool = await this.#ensurePool(opts.modelId); + const result = unwrap( + await pool.submit({ + ...scheduling(opts), + id: opts.id, + kind: JOB_KIND.EMBEDDING, + params: { input: texts }, + }), + ); + return { data: result.embeddings ?? [], usage: result.usage }; + } + + /** + * Independent prompts, fanned across the pool. This is the only way to beat + * the ~10 tok/s single-stream ceiling, so anything embarrassingly parallel + * (translating a page, labelling a list) should arrive here rather than as a + * loop of `complete` calls. + * + * @param {CompletionRequest & {requests: Array>}} payload + * @param {(item: BatchItem) => void} [onItem] called as each item lands + * @returns {Promise>} + */ + async batch(payload, onItem = () => {}) { + const requests = payload.requests; + if (!Array.isArray(requests) || requests.length === 0) { + throw new EngineError(ERROR.BAD_REQUEST, "`requests` must be a non-empty array."); + } + const pool = await this.#ensurePool(payload.modelId); + const sched = scheduling(payload); + // One batch is one task, however many requests it is: "translate this page" + // should hold one engine, not every engine. The pool reserves its last free + // slot for a different task, so ghost-text never queues behind the page. + const task = payload.task ?? `batch-${payload.id ?? crypto.randomUUID()}`; + + return Promise.all( + requests.map(async (req, index) => { + const merged = { ...payload, ...req, requests: undefined }; + const result = await pool.submit({ + ...sched, + ...scheduling(merged), + session: req.session, // a batch shares no session unless an item names one + // An item that names its own session is its own task again. + task: req.task ?? req.session ?? task, + params: await this.#buildParams(merged), + }); + const item = { + index, + engineIndex: result.engineIndex, + startedAt: result.startedAt, + finishedAt: result.finishedAt, + ...(result.error + ? { error: result.error } + : { + text: result.text, + usage: result.usage, + finishReason: result.finishReason, + ...(result.toolCalls ? { toolCalls: result.toolCalls } : {}), + ...(result.cancelled ? { cancelled: true } : {}), + }), + }; + onItem(item); + return item; + }), + ); + } + + /** + * Cancels by job id or by session key. + * @param {string} idOrSession + * @returns {number} how many jobs it stopped + */ + cancel(idOrSession) { + let stopped = 0; + for (const pool of this.#pools.values()) stopped += pool.cancel(idOrSession); + return stopped; + } + + /** + * Applies a runtime knob to the running pool and persists it as the default. + * + * `decodeSteps` is the multi-step decode width (AI.md, "Multi-step decoding"). + * It takes effect on the next burst — no reload — which is what makes sweeping + * it to find this machine's tick boundary cheap. + */ + async configure(patch) { + const applied = {}; + if (patch.decodeSteps !== undefined) applied.decodeSteps = clampSteps(patch.decodeSteps); + if (patch.engineCount !== undefined) { + const n = Math.round(Number(patch.engineCount)); + if (!Number.isFinite(n) || n < 1) { + throw new EngineError( + ERROR.BAD_REQUEST, + `engineCount must be a positive integer, not ${JSON.stringify(patch.engineCount)}.`, + { engineCount: patch.engineCount }, + ); + } + applied.engineCount = n; + } + if (Object.keys(applied).length === 0) { + // Naming the knobs matters: this is the error a caller hits after + // `environment()` told them something was operable, so it has to agree + // with that report about what the operable things are. + throw new EngineError( + ERROR.BAD_REQUEST, + "`configure` needs at least one setting. Operable: `decodeSteps`, `engineCount`.", + { operable: ["decodeSteps", "engineCount"] }, + ); + } + await this.#store.setSettings(applied); + // Only `decodeSteps` is hot. `engineCount` is persisted and read when a pool + // is built, so a live pool keeps the size it came up with — `environment()` + // reports that gap rather than pretending the change took effect. + let engines = 0; + if (applied.decodeSteps !== undefined) { + for (const pool of this.#pools.values()) engines += pool.configure({ decodeSteps: applied.decodeSteps }); + } + return { settings: applied, engines }; + } +} + +/** + * Normalise every folder shape a caller might hold into `{path, file}[]`. + * + * A drop event gives a `DataTransfer`, `` gives a + * `FileList`, and a caller who has already unpacked one gives the entries. All + * three mean "this folder", so `load()` accepts all three rather than making + * the caller find the right converter first. + */ +async function toEntries(files) { + if (isDataTransfer(files)) return filesFromDataTransfer(files); + if (isFileList(files)) return filesFromInput(files); + return files; +} + +/** Scheduling metadata is per-request; the pool, not the caller, acts on it. */ +function scheduling(payload) { + return { + task: payload.task, + session: payload.session, + priority: payload.priority ?? PRIORITY.NORMAL, + preemptible: payload.preemptible, + }; +} + +function unwrap(result) { + // The pool resolves rather than rejects, so a failed generation arrives as a + // string on the result. It has no code of its own by then. + if (result.error) throw asEngineError(result.error, ERROR.GENERATION_FAILED); + return result; +} diff --git a/src/engine/environment.js b/src/engine/environment.js new file mode 100644 index 0000000000000000000000000000000000000000..ff72bcc2449e1e763e1d425d2508204e94c20185 --- /dev/null +++ b/src/engine/environment.js @@ -0,0 +1,295 @@ +/** + * `environment()` — one call that says what this machine will do, and why. + * + * It absorbs the *read* half of `probe()`, `features()` and `estimateSpeed()`. + * It is **read-only**. Writes go through `configure()`, which already existed + * and stays the one way to change a setting. + * + * That split was a deliberate correction to the original plan, which had one + * function do both and told them apart by argument shape. Implicit dispatch is + * the opposite of foolproof: `environment({ scope: "device", decodeSteps: 8 })` + * has no intuitive answer, and the reason it has none is that one function was + * being asked to do two jobs. So reading and writing are different verbs, and + * the report's `fix` field names the write to make instead of performing it. + * + * Every line carries the same five fields, because a caller should be able to + * render the whole report without special-casing any of it: + * + * severity blocked | degraded | tune | info | ok + * affects what the reader loses — in their terms, not the engine's + * cause the measured fact behind the verdict + * fix the exact call or instruction, or `null` when nothing can be done + * operable whether this is reachable from JS at all + * + * `fix` is `null` for anything the caller genuinely cannot change — hardware, + * build-time flags, browser settings that JS cannot reach. Reporting a + * consequence with no remedy is still worth doing: "your second turn is slow + * because this device caps storage buffers at 9" is the difference between a + * bug report and an informed decision. + * + * **`canRun(modelId)` stays where it is.** It is per-*model*; this is + * per-*device*, and model ranking belongs with model discovery. The two answer + * different questions and merging them would make both worse. + */ +import { ERROR, EngineError } from "./errors.js"; + +export const SEVERITY = { + /** Nothing will run until this is dealt with. */ + BLOCKED: "blocked", + /** It runs, measurably worse. */ + DEGRADED: "degraded", + /** It runs well; this is a dial worth turning. */ + TUNE: "tune", + /** Worth knowing, nothing to do. */ + INFO: "info", + /** Checked, and fine. */ + OK: "ok", +}; + +/** Worst-first, so `ok` can be computed and lines sorted by urgency. */ +const RANK = [SEVERITY.BLOCKED, SEVERITY.DEGRADED, SEVERITY.TUNE, SEVERITY.INFO, SEVERITY.OK]; + +/** Paged prefill's binding count; below this, cross-turn KV reuse cannot build. */ +const PAGED_PREFILL_STORAGE_BUFFERS = 10; + +const SCOPES = new Set(["full", "local", "device"]); + +/** + * Builds the callable `engine.environment` — a function with `.measure()` on it. + * + * A function rather than an object because the common case is asking for the + * whole report, and `environment()` should be the short spelling of that. + * + * @param {import("./engine.js").ScheduledEngine} engine + */ +export function environmentFacade(engine) { + const environment = (opts = {}) => report(engine, opts); + + /** + * One calibration generation, then the measured rate. + * + * Returns tok/s for the **current model** — the only one it can measure, + * since measuring means generating. The device-level bytes/sec figure it + * teaches the engine is what projects *other* models, and that projection is + * `estimateSpeed(id)`'s job, not this one's. + */ + environment.measure = async ({ tokens = 32 } = {}) => { + const modelId = engine.state.modelId; + if (!modelId || engine.resident.length === 0) { + throw new EngineError( + ERROR.NO_MODEL, + "environment.measure() generates in order to measure, so it needs a resident model. " + + "Call load() first, or use estimateSpeed(id) for a projection.", + ); + } + await engine.complete({ + messages: [{ role: "user", content: "Count from one to twenty." }], + max_tokens: tokens, + // Calibration must not jump a queue a real caller is waiting in. + priority: "background", + preemptible: true, + }); + return { ...(await engine.estimateSpeed(modelId)), tokens }; + }; + + return environment; +} + +async function report(engine, { scope = "full", ...rest } = {}) { + // The one thing an implicit read/write API could not do: notice that a caller + // meant to write. `configure()` is named in the error because that is the + // call they wanted. + const stray = Object.keys(rest); + if (stray.length) { + throw new EngineError( + ERROR.BAD_REQUEST, + `environment() reports; it does not change anything. ` + + `To set ${stray.map((k) => `\`${k}\``).join(", ")}, call configure({ ${stray.join(", ")} }).`, + { keys: stray }, + ); + } + if (!SCOPES.has(scope)) { + throw new EngineError( + ERROR.BAD_REQUEST, + `environment() scope must be ${[...SCOPES].map((s) => `"${s}"`).join(", ")}, not "${scope}".`, + { scope }, + ); + } + + const device = await engine.probe(); + const lines = [...deviceLines(device)]; + + // A blocked device short-circuits everything below it. "K=15 forward steps + // per GPU sync" is true and completely useless next to "no model can load" — + // and burying the one line that matters under four that do not is exactly the + // failure this report exists to prevent. + const blocked = lines.some((l) => l.severity === SEVERITY.BLOCKED); + + if (scope !== "device" && !blocked) { + const features = await engine.features(); + const settings = await engine.store.getSettings(); + lines.push(...runtimeLines(engine, device, features, settings)); + } + + // `full` is the only scope that may touch the network: `estimateSpeed` reads + // the model list, which pulls the ~6 MB WebLLM bundle when prebuilt models + // are on. `local` exists precisely so a caller can ask cheaply and often. + let speed = null; + if (scope === "full" && !blocked && engine.state.modelId) { + speed = await engine.estimateSpeed().catch(() => null); + } + + lines.sort((a, b) => RANK.indexOf(a.severity) - RANK.indexOf(b.severity)); + const worst = lines.reduce( + (acc, l) => (RANK.indexOf(l.severity) < RANK.indexOf(acc) ? l.severity : acc), + SEVERITY.OK, + ); + + return { + scope, + ok: worst !== SEVERITY.BLOCKED, + severity: worst, + device, + ...(speed ? { speed } : {}), + lines, + }; +} + +const line = (id, severity, affects, cause, fix = null) => ({ + id, + severity, + affects, + cause, + fix, + operable: fix !== null, +}); + +function* deviceLines(device) { + if (!device.webgpu) { + // `reason` already carries the exact per-browser instruction, and it is not + // reachable from JS — so it is the fix text, not an operable one. + yield { + ...line("webgpu", SEVERITY.BLOCKED, "everything — no model can load", device.reason ?? "navigator.gpu is absent"), + fix: device.reason ?? null, + operable: false, + }; + return; + } + yield line("webgpu", SEVERITY.OK, null, "navigator.gpu is present and an adapter was granted"); + + const buffers = device.limits?.maxStorageBuffersPerShaderStage; + if (device.kvReuse === false) { + yield line( + "kvReuse", + SEVERITY.DEGRADED, + "every turn after the first re-prefills the whole history, so a long conversation waits seconds for its first token", + `paged prefill binds ${PAGED_PREFILL_STORAGE_BUFFERS} storage buffers per stage; this device allows ${buffers ?? "fewer"}`, + // Genuinely not operable: it is a driver/browser limit, not a setting. + null, + ); + } else if (device.kvReuse) { + yield line("kvReuse", SEVERITY.OK, null, `storage buffers per stage: ${buffers}`); + } + + yield device.features?.shaderF16 + ? line("shaderF16", SEVERITY.OK, null, "shader-f16 is supported, so q4f16 models run at full speed") + : line( + "shaderF16", + SEVERITY.DEGRADED, + "q4f16 models fall back to f32 maths, roughly halving decode", + "the adapter does not expose shader-f16", + null, + ); + + const { quota, usage, persisted } = device.storage ?? {}; + if (persisted === false) { + yield line( + "persist", + SEVERITY.TUNE, + "the browser may evict a multi-GB model under storage pressure, forcing a re-download", + "storage is not marked persistent", + "await ensurePersistent() — from everything-webgpu/adapters/idb. One-way, and may prompt.", + ); + } else if (persisted) { + yield line("persist", SEVERITY.OK, null, "storage is persistent; models will not be evicted"); + } + + if (quota && usage !== undefined) { + const freeGB = (quota - usage) / 1e9; + yield freeGB < 2 + ? line( + "quota", + SEVERITY.DEGRADED, + "a model may fail to cache, or evict one already there", + `${freeGB.toFixed(1)} GB free of ${(quota / 1e9).toFixed(1)} GB`, + "Free space, or remove(id) a model you no longer need.", + ) + : line("quota", SEVERITY.INFO, null, `${freeGB.toFixed(1)} GB free of ${(quota / 1e9).toFixed(1)} GB`); + } +} + +function* runtimeLines(engine, device, features, settings) { + // The 2d guard's report. Nothing else surfaces it, and its whole point is + // that the failure is otherwise silent. + const off = engine.state.decode?.multiStepOff; + if (off?.length) { + yield line( + "multiStepDecoding", + SEVERITY.DEGRADED, + "decode fell back to one GPU sync per token — roughly half throughput", + `the live pipeline is missing ${off.length} tvmjs internal(s): ${off.slice(0, 3).join("; ")}`, + "This is what a WebLLM upgrade looks like. Run `npm test` (webllm-contract) to see whether the names are gone from the bundle too.", + ); + } else if (features.multiStepDecoding) { + yield line("multiStepDecoding", SEVERITY.OK, null, `K=${features.decodeSteps} forward steps per GPU sync`); + } + + yield features.decodeSteps === 1 + ? line( + "decodeSteps", + SEVERITY.TUNE, + "decode pays one GPU sync per token, which is the ~10 tok/s ceiling", + "decodeSteps is 1, so multi-step decoding is off", + "configure({ decodeSteps: 15 })", + ) + : line( + "decodeSteps", + SEVERITY.INFO, + null, + `K=${features.decodeSteps}; the best value falls as the model grows, since it is ms/step that fills the 100 ms tick`, + "configure({ decodeSteps: n })", + ); + + yield line( + "engineCount", + SEVERITY.INFO, + null, + `pool cap ${features.maxEngines}, ${features.engines} built; it grows only when a second task waits`, + "configure({ engineCount: n }) — persisted, and applies to pools built after it", + ); + + // Observed, not asserted: a build with NO_PASS_MERGE=1 reports ~1. + const batching = features.computePassBatching; + if (batching !== null) { + yield batching > 1.5 + ? line("computePassBatching", SEVERITY.OK, null, `${batching.toFixed(1)} kernel launches per flush`) + : line( + "computePassBatching", + SEVERITY.DEGRADED, + "one compute pass per kernel launch, which measured ~2.5x slower", + `${batching.toFixed(1)} launches per flush — the build-time patch is not in effect`, + // Build-time, not runtime: nothing a caller holding the engine can do. + null, + ); + } + + if (settings.engineCount && features.maxEngines && settings.engineCount !== features.maxEngines) { + yield line( + "engineCountPending", + SEVERITY.INFO, + null, + `engineCount is ${settings.engineCount} but the live pool was built with ${features.maxEngines}`, + "Reload the model to apply it.", + ); + } +} diff --git a/src/engine/errors.js b/src/engine/errors.js new file mode 100644 index 0000000000000000000000000000000000000000..37a35f716a59c4244029c8589e88ed869e663bb0 --- /dev/null +++ b/src/engine/errors.js @@ -0,0 +1,85 @@ +/** + * Typed failures. + * + * Every error the engine raised used to be a bare string, so a caller wanting + * to tell "this device has no WebGPU" from "that folder was missing a shard" + * had to match on prose — which then silently broke whenever the prose + * improved. A code is the part of an error message that is allowed to be + * depended on. + * + * The codes are deliberately few. Each one exists because a caller does + * something *different* about it, not because it names a different line of + * code: + * + * NO_WEBGPU tell the user to check flags/hardware; retrying is futile + * NO_MODEL nothing registered at all — send them to your setup flow + * UNKNOWN_MODEL that id is not resolvable; `listAvailableModels()` says what is + * CACHE_INCOMPLETE a locally-registered model was evicted; re-register the folder + * INVALID_MODEL_FOLDER the folder is not a compiled MLC model; `detail` says what is missing + * BAD_REQUEST the caller's arguments are wrong; a bug in the caller + * ABORTED the caller cancelled it; not a failure, and not to be + * reported to a user as one + * GENERATION_FAILED the model failed mid-generation + * PACKAGE_INCOMPLETE this package is not wired into your build correctly — + * a missing `vendor/` bundle, or a decode worker the + * bundler did not emit. Fixed once, in the build, by + * the developer; never branched on at runtime. + * + * `PACKAGE_INCOMPLETE` covers two causes with one code on purpose. They are + * different sentences and `detail.cause` tells them apart, but no caller writes + * a different `catch` branch for them: both mean "your build is wrong, the app + * has not shipped yet", and both are fixed by editing config, not by handling + * an error path. A second code would have grown the table a caller reads + * without giving that caller anything new to do. + * + * `message` stays human-readable and stays the thing you print. `detail` + * carries whatever structured context the site had — the missing cache keys, + * the absent field — so a caller never has to parse the sentence. + */ +export const ERROR = { + NO_WEBGPU: "NO_WEBGPU", + NO_MODEL: "NO_MODEL", + UNKNOWN_MODEL: "UNKNOWN_MODEL", + CACHE_INCOMPLETE: "CACHE_INCOMPLETE", + INVALID_MODEL_FOLDER: "INVALID_MODEL_FOLDER", + BAD_REQUEST: "BAD_REQUEST", + ABORTED: "ABORTED", + GENERATION_FAILED: "GENERATION_FAILED", + PACKAGE_INCOMPLETE: "PACKAGE_INCOMPLETE", +}; + +export class EngineError extends Error { + /** + * @param {string} code one of ERROR + * @param {string} message human-readable, safe to show a user + * @param {object} [detail] structured context, so callers need not parse `message` + */ + constructor(code, message, detail) { + super(message); + this.name = "EngineError"; + this.code = code; + if (detail !== undefined) this.detail = detail; + } + + /** Wire form. `error` stays a plain string so existing callers keep working. */ + toJSON() { + return { code: this.code, message: this.message, ...(this.detail ? { detail: this.detail } : {}) }; + } +} + +/** + * Normalise anything thrown into an EngineError. + * + * Errors from WebLLM, the GPU stack and the structured-clone boundary arrive as + * plain Errors, DOMExceptions or strings, and a caller should not have to care + * which. Anything unrecognised becomes `GENERATION_FAILED` rather than being + * given a more specific code it has not earned. + */ +export function asEngineError(err, fallback = ERROR.GENERATION_FAILED) { + if (err instanceof EngineError) return err; + const message = String(err?.message ?? err); + return new EngineError(fallback, message); +} + +export const isEngineError = (err, code) => + err instanceof EngineError && (code === undefined || err.code === code); diff --git a/src/engine/index.js b/src/engine/index.js new file mode 100644 index 0000000000000000000000000000000000000000..197deb481021f975b3651c767e9696d48d7426d9 --- /dev/null +++ b/src/engine/index.js @@ -0,0 +1,89 @@ +/** + * Public entry point. + * + * Nothing reachable from here touches `browser.*`. A page, a worker, an + * extension background page and a Node test all import the same modules; only + * the StorageAdapter differs (`src/adapters/`). + * + * Migrating off `@mlc-ai/web-llm` is one line — the call below is unchanged: + * + * ```js + * import { CreateScheduledEngine } from "everything-webgpu"; + * + * const engine = await CreateScheduledEngine("Llama-3.2-1B-Instruct-q4f16_1-MLC"); + * const reply = await engine.chat.completions.create({ + * messages: [{ role: "user", content: "hi" }], + * session: "ghost-text", // added by this engine + * priority: "interactive", // added by this engine + * }); + * ``` + * + * The longer form, when you want to choose the store or the model source: + * + * ```js + * import { ScheduledEngine, ModelStore, ingestModelFolder } from "everything-webgpu"; + * import { indexedDBStorage } from "everything-webgpu/adapters/idb"; + * + * const engine = new ScheduledEngine({ store: new ModelStore(await indexedDBStorage()) }); + * + * // Weights arrive by any of three routes. One call covers the latter two. + * await engine.load("Llama-3.2-1B-Instruct-q4f16_1-MLC"); // prebuilt, from HuggingFace + * + * await engine.registerModel({ // a base URL you host + * modelId: "my-model", + * model: "/models/my-model/", + * modelLib: "/models/my-model/my-model-webgpu.wasm", + * }); + * await engine.registerModel({ modelId: "my-model", files }); // off disk, never any network + * + * const { text } = await engine.complete({ messages: [{ role: "user", content: "hi" }] }); + * ``` + * + * `listAvailableModels()` enumerates all three. Pass `{ prebuilt: false }` to + * the engine for a build that must never fetch a model over the network. + * + * The scheduling fields — `task`, `session`, `priority`, `preemptible` — are + * what this adds over calling WebLLM directly. See AI.md, "The three shapes of + * work". + */ +export { ScheduledEngine } from "./engine.js"; +export { CreateScheduledEngine } from "./create.js"; +export { + canRun, + probeDevice, + projectSpeed, + rankModels, + REFERENCE_DECODE_BYTES_PER_SECOND, +} from "./device.js"; +export { ERROR, EngineError, asEngineError, isEngineError } from "./errors.js"; +export { EnginePool } from "./pool.js"; + +export { + ModelStore, + MODEL_TYPE, + SOURCE, + isInjected, + DEFAULT_SETTINGS, + CACHE_CONFIG, + CACHE_MODEL, + CACHE_WASM, + baseUrlFor, + groupKeysByScope, + toAppConfig, + formatBytes, +} from "./model-store.js"; + +export { filesFromDataTransfer, filesFromInput, ingestModelFolder } from "./ingest.js"; +export { prefetchModel, resolveModelUrl } from "./prefetch.js"; +export { ask, conversation, ghostText } from "./recipes.js"; + +export { + DEFAULT_DECODE_STEPS, + MAX_DECODE_STEPS, + burstSize, + clampSteps, + installMultiStepDecoding, +} from "./multistep.js"; + +export { ENGINE_STATE, PRIORITY, PRIORITY_ORDER, UNLOAD_LEVEL, WORKER_CONFIGURE } from "./constants.js"; +export { SEVERITY } from "./environment.js"; diff --git a/src/engine/ingest.js b/src/engine/ingest.js new file mode 100644 index 0000000000000000000000000000000000000000..c51c53a680ecacc08e05cdbd16c9eecb54105d52 --- /dev/null +++ b/src/engine/ingest.js @@ -0,0 +1,294 @@ +/** + * Drag-and-drop ingestion: turn a folder of MLC-compiled model artifacts into + * populated Cache Storage entries, with no network involved at any point. + * + * Everything is validated before the first byte is written, so a folder that is + * missing a shard fails immediately instead of after copying 2 GB. + */ +import { ERROR, EngineError } from "./errors.js"; +import { CACHE_CONFIG, CACHE_MODEL, CACHE_WASM, SOURCE, baseUrlFor, toModelType } from "./model-store.js"; + +/** WebLLM asks for this name; older MLC exports ship `ndarray-cache.json`. */ +export const TENSOR_MANIFEST = "tensor-cache.json"; +export const LEGACY_TENSOR_MANIFEST = "ndarray-cache.json"; +export const CHAT_CONFIG = "mlc-chat-config.json"; + +export const CONTENT_TYPES = { + json: "application/json", + wasm: "application/wasm", + bin: "application/octet-stream", +}; + +/** + * Walk a DataTransfer from a drop event into flat `{ path, file }` entries. + * Uses the entries API so dropping a *folder* works, not just a file selection. + */ +export async function filesFromDataTransfer(dataTransfer) { + // `Array.from` throughout, for the reason `filesFromInput` gives: a + // DataTransferItemList and a FileList are array-like, and only sometimes + // iterable. Spreading them threw from inside here, three frames from the drop + // handler the caller actually wrote. + const roots = Array.from(dataTransfer.items) + .filter((item) => item.kind === "file") + .map((item) => (item.webkitGetAsEntry ? item.webkitGetAsEntry() : null)); + + if (roots.some((entry) => entry === null)) { + // No entries API: fall back to the flat file list (a folder drop yields nothing). + return Array.from(dataTransfer.files, (file) => ({ + path: file.webkitRelativePath || file.name, + file, + })); + } + + const out = []; + await Promise.all(roots.filter(Boolean).map((entry) => walkEntry(entry, "", out))); + return out; +} + +async function walkEntry(entry, prefix, out) { + const path = prefix ? `${prefix}/${entry.name}` : entry.name; + if (entry.isFile) { + out.push({ path, file: await new Promise((res, rej) => entry.file(res, rej)) }); + return; + } + const reader = entry.createReader(); + // readEntries() returns at most ~100 entries per call; drain it. + for (;;) { + const batch = await new Promise((res, rej) => reader.readEntries(res, rej)); + if (batch.length === 0) break; + await Promise.all(batch.map((child) => walkEntry(child, path, out))); + } +} + +/** + * Turns `` output into the same `{ path, file }` shape. + * + * `Array.from`, not spread: a real `FileList` is iterable, but plenty of things + * that behave like one are only array-like, and spreading those fails with + * "fileList is not iterable" — an error that names none of the three places it + * could have come from. Array.from accepts both. + */ +export function filesFromInput(fileList) { + return Array.from(fileList, (file) => ({ + path: file.webkitRelativePath || file.name, + file, + })); +} + +/** + * @param {Array<{path: string, file: File}>} entries + * @param {object} opts + * @param {import("./model-store.js").ModelStore} opts.store where the registry entry lands + * @param {string} [opts.modelId] overrides the id inferred from the folder name + * @param {"llm"|"embedding"|"vlm"} [opts.modelType] declare a vision model, or WebLLM + * rejects every image sent to it + * @param {(p: {phase: string, done: number, total: number, label: string}) => void} [opts.onProgress] + * @returns {Promise} the saved registry record + */ +export async function ingestModelFolder(entries, { store, modelId, modelType, onProgress = () => {} } = {}) { + if (!store) { + throw new EngineError(ERROR.BAD_REQUEST, "ingestModelFolder needs a `store` to save the registry entry into."); + } + if (!entries?.length) { + throw new EngineError(ERROR.INVALID_MODEL_FOLDER, "Nothing was dropped — expected a model folder.", { + reason: "empty", + }); + } + + const byPath = new Map(); + const byName = new Map(); + for (const { path, file } of entries) { + const relative = stripRoot(path); + byPath.set(relative, file); + // Last writer wins; `find()` prefers the exact relative path anyway. + byName.set(basename(relative), file); + } + + const find = (name) => byPath.get(name) ?? byName.get(basename(name)); + + onProgress({ phase: "validating", done: 0, total: 1, label: "Reading manifests" }); + + const configFile = find(CHAT_CONFIG); + if (!configFile) { + throw new EngineError( + ERROR.INVALID_MODEL_FOLDER, + `Missing ${CHAT_CONFIG}. Use the folder produced by \`mlc_llm convert_weights\` + \`gen_config\`, not a raw HuggingFace checkpoint.`, + { reason: "missing-config", missing: [CHAT_CONFIG] }, + ); + } + const chatConfig = await readJson(configFile, CHAT_CONFIG); + + const tensorFile = find(TENSOR_MANIFEST) ?? find(LEGACY_TENSOR_MANIFEST); + if (!tensorFile) { + throw new EngineError( + ERROR.INVALID_MODEL_FOLDER, + `Missing ${TENSOR_MANIFEST} (or legacy ${LEGACY_TENSOR_MANIFEST}) — the weight shard index.`, + { reason: "missing-manifest", missing: [TENSOR_MANIFEST] }, + ); + } + const isLegacyManifest = !find(TENSOR_MANIFEST); + const tensorManifest = await readJson(tensorFile, tensorFile.name); + + const records = tensorManifest.records; + if (!Array.isArray(records) || records.length === 0) { + throw new EngineError( + ERROR.INVALID_MODEL_FOLDER, + `${tensorFile.name} has no "records" array — it is not an MLC weight index.`, + { reason: "malformed-manifest" }, + ); + } + + const shardPaths = records.map((r) => r.dataPath).filter(Boolean); + if (shardPaths.length !== records.length) { + throw new EngineError( + ERROR.INVALID_MODEL_FOLDER, + `${tensorFile.name} has records without a "dataPath".`, + { reason: "malformed-manifest" }, + ); + } + const missingShards = shardPaths.filter((p) => !find(p)); + if (missingShards.length) { + throw new EngineError( + ERROR.INVALID_MODEL_FOLDER, + `${missingShards.length} weight shard(s) missing from the folder: ${missingShards.slice(0, 5).join(", ")}${missingShards.length > 5 ? ", …" : ""}`, + { reason: "missing-shards", missing: missingShards }, + ); + } + + const tokenizerNames = Array.isArray(chatConfig.tokenizer_files) ? chatConfig.tokenizer_files : []; + const tokenizerName = ["tokenizer.json", "tokenizer.model"].find( + (name) => tokenizerNames.includes(name) && find(name), + ); + if (!tokenizerName) { + throw new EngineError( + ERROR.INVALID_MODEL_FOLDER, + `No usable tokenizer. ${CHAT_CONFIG} lists [${tokenizerNames.join(", ") || "nothing"}], and neither tokenizer.json nor tokenizer.model is present in the folder.`, + { reason: "missing-tokenizer", expected: tokenizerNames }, + ); + } + + const wasmEntries = [...byPath.entries()].filter(([p]) => p.endsWith(".wasm")); + if (wasmEntries.length === 0) { + throw new EngineError( + ERROR.INVALID_MODEL_FOLDER, + "No .wasm model library found. Add the matching `*-webgpu.wasm` from mlc-ai/binary-mlc-llm-libs to the folder.", + { reason: "missing-wasm" }, + ); + } + if (wasmEntries.length > 1) { + throw new EngineError( + ERROR.INVALID_MODEL_FOLDER, + `Found ${wasmEntries.length} .wasm files (${wasmEntries.map(([p]) => p).join(", ")}); the folder must contain exactly one model library.`, + { reason: "ambiguous-wasm", found: wasmEntries.map(([p]) => p) }, + ); + } + const [wasmPath, wasmFile] = wasmEntries[0]; + const wasmName = basename(wasmPath); + + const id = (modelId || inferModelId(entries) || wasmName.replace(/(-webgpu)?\.wasm$/, "")).trim(); + if (!id) { + throw new EngineError( + ERROR.INVALID_MODEL_FOLDER, + "Could not determine a model id — name the folder after the model, or pass `modelId`.", + { reason: "no-model-id" }, + ); + } + + const base = baseUrlFor(id); + + // Everything validated: build the write plan. + const plan = [ + { scope: CACHE_CONFIG, url: base + CHAT_CONFIG, file: configFile, type: CONTENT_TYPES.json }, + { scope: CACHE_MODEL, url: base + TENSOR_MANIFEST, file: tensorFile, type: CONTENT_TYPES.json }, + { + scope: CACHE_MODEL, + url: base + tokenizerName, + file: find(tokenizerName), + type: tokenizerName.endsWith(".json") ? CONTENT_TYPES.json : CONTENT_TYPES.bin, + }, + { scope: CACHE_WASM, url: base + wasmName, file: wasmFile, type: CONTENT_TYPES.wasm }, + ...shardPaths.map((p) => ({ + scope: CACHE_MODEL, + url: new URL(p, base).href, + file: find(p), + type: CONTENT_TYPES.bin, + })), + ]; + if (isLegacyManifest) { + // Keep the original name addressable too, so a future WebLLM that reverts + // to `ndarray-cache.json` still hits cache. + plan.push({ + scope: CACHE_MODEL, + url: base + LEGACY_TENSOR_MANIFEST, + file: tensorFile, + type: CONTENT_TYPES.json, + }); + } + + const openCaches = new Map(); + let done = 0; + for (const item of plan) { + if (!openCaches.has(item.scope)) openCaches.set(item.scope, await caches.open(item.scope)); + onProgress({ phase: "writing", done, total: plan.length, label: basename(item.url) }); + await openCaches.get(item.scope).put( + new Request(item.url), + new Response(item.file, { status: 200, headers: { "Content-Type": item.type } }), + ); + done += 1; + } + onProgress({ phase: "writing", done, total: plan.length, label: "done" }); + + const keys = { [CACHE_CONFIG]: [], [CACHE_MODEL]: [], [CACHE_WASM]: [] }; + for (const item of plan) keys[item.scope].push(item.url); + + return store.save({ + model_id: id, + model: base, + model_lib: base + wasmName, + source: SOURCE.INJECTED, + ...(toModelType(modelType) !== undefined ? { model_type: toModelType(modelType) } : {}), + ...(chatConfig.context_window_size > 0 + ? { overrides: { context_window_size: chatConfig.context_window_size } } + : {}), + keys, + sizeBytes: plan.reduce((sum, item) => sum + item.file.size, 0), + fileCount: plan.length, + shardCount: shardPaths.length, + tokenizer: tokenizerName, + wasm: wasmName, + addedAt: new Date().toISOString(), + }); +} + +async function readJson(file, label) { + try { + return JSON.parse(await file.text()); + } catch (err) { + throw new EngineError(ERROR.INVALID_MODEL_FOLDER, `${label} is not valid JSON: ${err.message}`, { + reason: "malformed-json", + file: label, + }); + } +} + +/** Drops the dropped-folder name so `Qwen3-4B/tokenizer.json` keys as `tokenizer.json`. */ +function stripRoot(path) { + const parts = path.split("/"); + return parts.length > 1 ? parts.slice(1).join("/") : path; +} + +function basename(path) { + return path.split("/").pop(); +} + +/** + * A dropped *folder* gives every entry the same first path segment; a flat + * multi-file selection gives each entry a bare filename. Only the former names + * the model. + */ +function inferModelId(entries) { + const roots = new Set( + entries.filter((e) => e.path.includes("/")).map((e) => e.path.split("/")[0]), + ); + return roots.size === 1 ? [...roots][0] : ""; +} diff --git a/src/engine/model-store.js b/src/engine/model-store.js new file mode 100644 index 0000000000000000000000000000000000000000..7082c4aef31c93608788eac3c88f9678376e9827 --- /dev/null +++ b/src/engine/model-store.js @@ -0,0 +1,371 @@ +/** + * The model registry, and the Cache Storage layout WebLLM expects. + * + * A model reaches the engine by one of three routes. They differ only in where + * the weights are fetched from; everything downstream is identical, because all + * three end up as one `model_list` entry WebLLM's own loader resolves. + * + * prebuilt one of the 163 entries in WebLLM's `prebuiltAppConfig`, on + * HuggingFace. Nothing to register: `load("Llama-3.2-1B-Instruct- + * q4f16_1-MLC")` just works. + * remote `registerModel({ model, modelLib })` with any base URL — an HF + * repo, your own CDN, a path on your own origin, localhost. This + * is how a developer points the engine at weights they host. + * injected `ingestModelFolder()` writes a local folder straight into Cache + * Storage. No network at any point, for offline or private builds. + * + * The injected route mints a synthetic https base URL and pre-populates the + * exact cache scopes/keys WebLLM's loader would have populated from the + * network, so `reload()` finds every artifact already cached and issues zero + * requests. That origin is `.invalid` on purpose: it can never resolve, so an + * injected model whose cache was evicted fails loudly instead of quietly + * pulling a gigabyte off the network. Enabling downloads for the other two + * routes cannot weaken that guarantee. + * + * Registry and settings hang off a `ModelStore` holding an injected + * `StorageAdapter`. That is the whole reason this file is no longer + * extension-bound: `browser.storage.local` was the only WebExtension API in the + * engine core outside the router. + * + * Cache Storage is deliberately *not* injected. `caches` exists in every secure + * context, and the cache keys are the contract with WebLLM's loader — putting + * an abstraction over them would hide the one thing that has to stay exact. + * + * Cache scopes (must stay in sync with @mlc-ai/web-llm): + * webllm/config -> mlc-chat-config.json + * webllm/model -> tensor-cache.json, tokenizer file, every shard + * webllm/wasm -> .wasm + */ +import { ERROR, EngineError } from "./errors.js"; + +export const CACHE_CONFIG = "webllm/config"; +export const CACHE_MODEL = "webllm/model"; +export const CACHE_WASM = "webllm/wasm"; + +/** + * `.invalid` is reserved by RFC 6761 and can never resolve, so a bug that skips + * the cache surfaces as a DNS failure instead of a silent download. + * The `/resolve/main/` suffix makes WebLLM's `cleanModelUrl()` a no-op. + */ +const VIRTUAL_ORIGIN = "https://local-model.invalid"; + +const STORAGE_KEY = "models"; +const SETTINGS_KEY = "settings"; + +export const DEFAULT_SETTINGS = { + /** Empty list = every installed extension may call the API. Wire adapter only. */ + allowedExternalIds: [], + /** + * Engines held in the pool. Each is a full copy of the weights in VRAM and a + * full load, but concurrent generations each get their own ~10 tok/s, so this + * is the only dial that raises total throughput. 2 is the smallest number + * that delivers any parallelism at all. + */ + engineCount: 2, + /** + * Forward steps per GPU->CPU sync (vLLM's `--num-scheduler-steps`). Decode is + * sync-bound, not compute-bound, so this is the only dial that raises + * *single-stream* throughput — `engineCount` raises aggregate throughput. + * + * 15 is vLLM's documented cap and this engine's default. Unlike vLLM the + * win here is quantized by Firefox's 100 ms poll, so the best value is the + * largest K whose burst still fits inside one tick, and it shrinks as the + * model grows. See src/engine/multistep.js and `npm run e2e -- --steps`. + */ + decodeSteps: 15, + /** + * `buildParams` puts this on every request, so it shadows whatever + * `mlc-chat-config.json` ships as the model's own default — unlike `top_p`, + * which is never injected and so comes from the model. 0.6 is what the + * Qwen3.8-2B-Distill card asks for; reasoning models in this class are prone + * to repetition loops when decoding is too close to greedy. + */ + temperature: 0.6, + maxTokens: 1024, + systemPrompt: "", +}; + +/** + * Make a base URL absolute, at registration rather than at load. + * + * WebLLM's `cleanModelUrl` ends in `new URL(url)` with no base, so it throws on + * a relative path — `/models/my-model/` fails deep inside the loader, long + * after the caller could tell why. Resolving here means a relative path works + * as documented, and a context with no page URL to resolve against says so + * immediately instead of at load time. + */ +function absolutize(url, field) { + if (/^[a-z][a-z0-9+.-]*:/i.test(url)) return url; + const base = globalThis.location?.href; + if (!base) { + throw new EngineError( + ERROR.BAD_REQUEST, + `\`${field}\` is relative ("${url}") and this context has no page URL to resolve it against. ` + + "Pass an absolute URL.", + { field, value: url }, + ); + } + return new URL(url, base).href; +} + +export function baseUrlFor(modelId) { + return `${VIRTUAL_ORIGIN}/${encodeURIComponent(modelId)}/resolve/main/`; +} + +/** Every cache key a record claims, keyed by cache scope. */ +export function groupKeysByScope(record) { + return { + [CACHE_CONFIG]: record.keys?.[CACHE_CONFIG] ?? [], + [CACHE_MODEL]: record.keys?.[CACHE_MODEL] ?? [], + [CACHE_WASM]: record.keys?.[CACHE_WASM] ?? [], + }; +} + +/** + * WebLLM's own `ModelType` enum, which it reads off the `model_list` entry. + * + * This matters for one reason: WebLLM refuses image content on anything not + * marked `VLM` (`UserMessageContentErrorForNonVLM`). It cannot be inferred — + * `mlc-chat-config.json` carries the architecture name, not this — so a + * locally compiled vision model has to declare it or every image is rejected + * with a confusing error. + */ +export const MODEL_TYPE = { llm: 0, embedding: 1, vlm: 2 }; + +/** Accepts `"vlm"`, `MODEL_TYPE.vlm`, or nothing. */ +export function toModelType(value) { + if (value === undefined || value === null) return undefined; + if (typeof value === "number") return value; + const known = MODEL_TYPE[String(value).toLowerCase()]; + if (known === undefined) { + throw new EngineError( + ERROR.BAD_REQUEST, + `Unknown modelType "${value}". Expected one of: ${Object.keys(MODEL_TYPE).join(", ")}.`, + { modelType: value }, + ); + } + return known; +} + +/** How a record's weights are obtained. See the header. */ +export const SOURCE = { + PREBUILT: "prebuilt", + REMOTE: "remote", + INJECTED: "injected", +}; + +/** + * Whether this record's bytes live in Cache Storage and nowhere else. + * + * The distinction that matters: an injected model that loses its cache is + * unrecoverable and must be re-ingested, so `verify()` gates its load. A remote + * or prebuilt one just re-downloads, so eviction is a slow load, not an error. + */ +export const isInjected = (record) => record?.source === SOURCE.INJECTED; + +function toModelListEntry(record) { + return { + model: record.model, + model_id: record.model_id, + model_lib: record.model_lib, + // Carried through, or WebLLM treats a locally registered VLM as text-only. + ...(record.model_type !== undefined ? { model_type: record.model_type } : {}), + ...(record.overrides ? { overrides: record.overrides } : {}), + ...(record.vram_required_MB ? { vram_required_MB: record.vram_required_MB } : {}), + }; +} + +/** + * Shape WebLLM's `appConfig` from the registry, optionally over its own + * prebuilt list. + * + * Registered records win on a model_id collision, so a developer can shadow a + * prebuilt entry — point `Llama-3.2-1B-Instruct-q4f16_1-MLC` at their own + * mirror, say — without renaming it and breaking their callers. + * + * @param {Array} models registered records + * @param {{model_list: Array} | null} [prebuilt] WebLLM's `prebuiltAppConfig` + */ +export function toAppConfig(models, prebuilt = null) { + const own = models.map(toModelListEntry); + const owned = new Set(own.map((e) => e.model_id)); + const rest = (prebuilt?.model_list ?? []).filter((e) => !owned.has(e.model_id)); + return { model_list: [...own, ...rest], useIndexedDBCache: false }; +} + +export function formatBytes(bytes) { + if (!Number.isFinite(bytes) || bytes <= 0) return "0 B"; + const units = ["B", "KB", "MB", "GB", "TB"]; + const i = Math.min(Math.floor(Math.log(bytes) / Math.log(1024)), units.length - 1); + return `${(bytes / 1024 ** i).toFixed(i === 0 ? 0 : 1)} ${units[i]}`; +} + +/** + * @typedef {object} StorageAdapter + * @property {(key: string | string[]) => Promise} get + * @property {(items: object) => Promise} set + * + * Two methods, deliberately the exact shape of `browser.storage.local`, so the + * WebExtension adapter is a passthrough rather than a translation layer. See + * `src/adapters/` for the three implementations. + */ + +export class ModelStore { + /** @type {StorageAdapter} */ + #storage; + + /** @param {StorageAdapter} storage */ + constructor(storage) { + if (!storage?.get || !storage?.set) { + throw new EngineError(ERROR.BAD_REQUEST, "ModelStore needs a StorageAdapter with `get(key)` and `set(obj)`."); + } + this.#storage = storage; + } + + /** @returns {Promise>} registered model records, newest first. */ + async list() { + const { [STORAGE_KEY]: models } = await this.#storage.get(STORAGE_KEY); + return Array.isArray(models) ? models : []; + } + + async get(modelId) { + return (await this.list()).find((m) => m.model_id === modelId); + } + + /** + * Register weights the engine should fetch rather than find already cached. + * + * This is the whole "developer configures a read path" surface: `model` is a + * base URL and nothing more, so an HF repo, your own CDN, a folder served off + * your own origin and a localhost dev server are all the same call. + * + * ```js + * await store.registerModel({ + * modelId: "Qwen3.8-2B-q4f16_1-MLC", + * model: "/models/Qwen3.8-2B-q4f16_1-MLC/", + * modelLib: "/models/Qwen3.8-2B-q4f16_1-MLC/Qwen3.8-2B-q4f16_1-webgpu.wasm", + * }); + * ``` + * + * No validation of the URLs happens here — there is nothing to validate + * without fetching, and WebLLM's loader reports a bad base URL far better + * than a HEAD request would. Contrast `ingestModelFolder`, which validates + * exhaustively because it is about to write a gigabyte. + */ + async registerModel({ modelId, model, modelLib, contextWindow, vramRequiredMB, modelType }) { + const missing = [ + ["modelId", modelId], + ["model", model], + ["modelLib", modelLib], + ].filter(([, v]) => !v); + if (missing.length) { + throw new EngineError( + ERROR.BAD_REQUEST, + `registerModel needs ${missing.map(([k]) => `\`${k}\``).join(", ")}. ` + + "Pass `files` instead to register a local folder.", + { missing: missing.map(([k]) => k) }, + ); + } + return this.save({ + model_id: modelId, + model: absolutize(model, "model"), + model_lib: absolutize(modelLib, "modelLib"), + source: SOURCE.REMOTE, + ...(toModelType(modelType) !== undefined ? { model_type: toModelType(modelType) } : {}), + ...(contextWindow > 0 ? { overrides: { context_window_size: contextWindow } } : {}), + ...(vramRequiredMB > 0 ? { vram_required_MB: vramRequiredMB } : {}), + addedAt: new Date().toISOString(), + }); + } + + async save(record) { + const models = (await this.list()).filter((m) => m.model_id !== record.model_id); + models.unshift(record); + await this.#storage.set({ [STORAGE_KEY]: models }); + return record; + } + + /** + * Frees an **injected** model's bytes; keeps the registry entry. + * + * Scoped to injected models deliberately. We wrote those artifacts, so we + * hold the only manifest of them — which also makes this stronger than + * WebLLM's equivalent *for this case*: `deleteModelAllInfoInCache` reads + * `tensor-cache.json` to enumerate shards, so it cannot clean up after that + * file has itself been evicted. An explicit key list has no such failure. + * + * Remote and prebuilt models are WebLLM's to delete — it fetched them and + * derives their keys exactly as its loader did. `ScheduledEngine.evict()` + * routes to whichever owns the model. + * + * @returns {Promise<{freedKeys: number}>} + */ + async evictInjected(modelId) { + const record = await this.get(modelId); + if (!record) throw new EngineError(ERROR.UNKNOWN_MODEL, `Unknown model "${modelId}"`, { modelId }); + + let freedKeys = 0; + for (const [scope, urls] of Object.entries(groupKeysByScope(record))) { + if (urls.length === 0) continue; + const cache = await caches.open(scope); + for (const url of urls) if (await cache.delete(url)) freedKeys += 1; + } + return { freedKeys }; + } + + /** + * The registry-only primitive: drops the entry, plus the cache keys an + * **injected** record lists. + * + * A remote record claims no keys, so this frees nothing for it — WebLLM + * downloaded those bytes and derives their names. Use + * `ScheduledEngine.remove()` for a full teardown; calling this directly on a + * remote model leaks its shards, because deleting the record also destroys + * the only record of where they came from. + */ + async remove(modelId) { + const record = await this.get(modelId); + if (!record) throw new EngineError(ERROR.UNKNOWN_MODEL, `Unknown model "${modelId}"`, { modelId }); + + for (const [scope, urls] of Object.entries(groupKeysByScope(record))) { + const cache = await caches.open(scope); + await Promise.all(urls.map((url) => cache.delete(url))); + } + + const models = (await this.list()).filter((m) => m.model_id !== modelId); + await this.#storage.set({ [STORAGE_KEY]: models }); + } + + /** + * Confirms the caches still hold everything the record promised. + * + * Storage eviction is silent, so this is what stands between a stale registry + * entry and WebLLM trying to fetch `local-model.invalid` mid-load. Inside a + * WebExtension with `unlimitedStorage` this was a defensive check; on an + * ordinary page origin, where a multi-GB model is evictable unless + * `navigator.storage.persist()` was granted, it is a core mechanism. + * + * @returns {Promise<{ok: boolean, missing: string[]}>} + */ + async verify(record) { + const missing = []; + for (const [scope, urls] of Object.entries(groupKeysByScope(record))) { + if (urls.length === 0) continue; + const cache = await caches.open(scope); + const present = new Set((await cache.keys()).map((req) => req.url)); + for (const url of urls) if (!present.has(url)) missing.push(url); + } + return { ok: missing.length === 0, missing }; + } + + async getSettings() { + const { [SETTINGS_KEY]: settings } = await this.#storage.get(SETTINGS_KEY); + return { ...DEFAULT_SETTINGS, ...(settings ?? {}) }; + } + + async setSettings(patch) { + const next = { ...(await this.getSettings()), ...patch }; + await this.#storage.set({ [SETTINGS_KEY]: next }); + return next; + } +} diff --git a/src/engine/multistep.js b/src/engine/multistep.js new file mode 100644 index 0000000000000000000000000000000000000000..34152080e911df6a50ef169eea0583b10840be67 --- /dev/null +++ b/src/engine/multistep.js @@ -0,0 +1,607 @@ +/** + * Multi-step decoding: N forward steps per GPU->CPU sync. + * + * Why this exists: decode here is not compute-bound, it is *sync*-bound. Firefox + * resolves `onSubmittedWorkDone()` / `mapAsync()` only on a 100 ms poll tick + * (AI.md, "The 10 tok/s ceiling"), and stock WebLLM needs exactly one sync per + * token — it reads the sampled token id back to JS before it can build the next + * step's input. One token per tick = 9.6 tok/s, of which ~7 ms is real compute. + * + * The fix is the one vLLM ships as `--num-scheduler-steps`: run K steps before + * paying the per-batch cost once. What makes it possible here without touching + * the compiled model is that WebLLM's sampling path is *already* on the GPU — + * `softmax_with_temperature`, `argsort_probs` and `sample_with_top_p` hand back + * an int32[1] device tensor, and `Tensor.copyFrom(Tensor)` is a device-to-device + * copy. So the sampled id feeds straight back into `embed` without ever becoming + * a JS number: + * + * embed -> decode -> penalties -> softmax -> argsort -> sample -> embed -> ... + * + * Each step stages its id into its own CPU tensor, and the burst ends with + * **one** `device.sync()`. tvmjs queues GPU->CPU copies into `pendingGPUToCPUCopy` + * and only awaits them in `sync()`, so K readbacks still cost one tick. + * + * Two things follow from the 100 ms grid, and they are why `steps` is a dial: + * + * - The win is quantized, not linear. A burst costs `ceil(K * perStepMs / 100)` + * ticks, so throughput is a sawtooth and the good values of K are the ones + * landing just under a boundary. On a 0.8B at ~7.3 ms/step that is K=13 + * (~130 tok/s); K=14 already spills into a second tick and halves it. + * - The best K shrinks as the model grows, because `perStepMs` grows. A model + * at 25 ms/step wants K=4, not K=15. + * + * Cost of the trick: the sampler cannot see its own output mid-burst. Repetition + * and presence/frequency penalties use the token history as it stood when the + * burst started, and stop conditions are only checked after the readback, so a + * burst can overshoot a stop token and must then be rewound. Both are the same + * trade vLLM makes. Anything needing per-token CPU feedback (grammar-constrained + * JSON, logprobs, a logit processor) falls back to single-step, where behaviour + * is identical to stock WebLLM. + * + * The other cost is that all of this drives ~30 undocumented tvmjs internals. A + * WebLLM upgrade that renames one does not break generation — it turns the fast + * path off and takes the throughput with it, silently. `PIPELINE_CONTRACT` below + * is that surface written down and checked against the live pipeline before the + * first burst, so the failure announces itself instead of being measured months + * later. + */ + +/** vLLM's documented sweet spot, and the value this extension ships. */ +export const DEFAULT_DECODE_STEPS = 15; + +/** + * Above this, the lookahead thrown away at a stop token outweighs the tick it + * saves, and the transient logits/argsort buffers stop being free. + */ +export const MAX_DECODE_STEPS = 32; + +export const clampSteps = (n) => Math.max(1, Math.min(MAX_DECODE_STEPS, Math.round(Number(n)) || 1)); + +// -------------------------------------------------- the pipeline contract ---- + +/** + * Every tvmjs pipeline internal a burst drives, and how each must behave. + * + * None of these are documented, none are part of WebLLM's public surface, and + * nothing upstream promises they will keep their names. The contract test checks + * them against the *bundle* on every `npm test`; this checks them against the + * *live object*, which is a different question — a member can survive in the + * bundle and still not be on the pipeline handed to us, if upstream moves it to + * a subclass, a different pipeline type, or behind a factory. + * + * Split three ways because presence alone is not the failure that hurts: + * + * - **`calls`** must be callable. A rename here throws, which is the *good* + * case — it is loud. + * - **`numbers`** are read arithmetically or incremented in place. This is the + * silent one: `pipeline.filledKVCacheLength += 1` on a member that no longer + * exists creates a new property, nothing throws, and the KV cache accounting + * quietly drifts. A missing `contextWindowSize` makes `burstSize` NaN. + * - **`reads`** need only exist. + * + * `logitProcessor` is deliberately optional: `burstSize` tests it for + * `undefined`, so absent is the normal case, not a broken one. + * + * The list is not maintained by hand — `webllm-contract.test.mjs` derives the + * set this file actually reaches for from its own source and asserts it matches + * this declaration exactly, so adding a `pipeline.newThing` without declaring it + * fails the build. + */ +export const PIPELINE_CONTRACT = { + calls: [ + "embed", + "fKVCacheBeginForward", + "fKVCacheEndForward", + "fapplyLogitBias", + "fapplyPenalty", + "fargsortProbs", + "fsampleWithTopP", + "fsoftmaxWithTemperature", + "getActiveKVStates", + "invokeDecode", + "processNextToken", + "resetChat", + "stopped", + ], + numbers: [ + "contextWindowSize", + "curRoundDecodingTotalTime", + "curRoundDecodingTotalTokens", + "decodingTotalTime", + "decodingTotalTokens", + "filledKVCacheLength", + "fullVocabSize", + "slidingWindowSize", + ], + reads: [ + "appearedTokensFreq", + "config", + "device", + "outputIds", + "params", + "sampleIndices", + "sampleIndicesDevice", + "topPDevice", + "tvm", + ], + optional: ["logitProcessor"], +}; + +/** + * What this pipeline is missing, as sentences a reader can act on. + * Empty means a burst is safe to run. + */ +export function missingPipelineMembers(pipeline) { + if (!pipeline || typeof pipeline !== "object") return ["the pipeline itself is not an object"]; + const missing = []; + for (const name of PIPELINE_CONTRACT.calls) { + if (typeof pipeline[name] !== "function") missing.push(`${name}() is not a function`); + } + for (const name of PIPELINE_CONTRACT.numbers) { + if (typeof pipeline[name] !== "number") missing.push(`${name} is not a number`); + } + for (const name of PIPELINE_CONTRACT.reads) { + if (pipeline[name] === undefined) missing.push(`${name} is missing`); + } + return missing; +} + +/** + * Replaces `engine.decode` with a burst-and-drain version. + * + * `decode` stays a one-token call — the caller's loop still checks + * `pipeline.stopped()` between tokens and still emits one chunk per token — but + * only one call in K actually touches the GPU. The rest drain a buffer. + * + * @param {object} engine an MLCEngine; in this project the one inside the worker + * @param {object} [options] + * @param {number} [options.steps] forward steps per sync; 1 disables the path + * @param {(info: {steps: number, tokens: number, ms: number}) => void} [options.onBurst] + * @param {(info: {missing: string[]}) => void} [options.onFallback] fired once + * per pipeline that fails the contract, before it is routed to stock decoding + * @returns {{setSteps: (n: number) => void, readonly steps: number, + * readonly fallbacks: number}} + */ +export function installMultiStepDecoding( + engine, + { steps = DEFAULT_DECODE_STEPS, onBurst, onFallback } = {}, +) { + const config = { steps: clampSteps(steps), fallbacks: 0 }; + const lookahead = new WeakMap(); + const baseDecode = engine.decode.bind(engine); + const basePrefill = engine.prefill.bind(engine); + + const stateFor = (pipeline) => { + let state = lookahead.get(pipeline); + if (!state) lookahead.set(pipeline, (state = { queue: [] })); + return state; + }; + + /** Contract verdict per pipeline; the check runs once, the answer is reused. */ + const supported = new WeakMap(); + + /** + * Whether this pipeline may be burst, decided once and remembered. + * + * Checked at first decode rather than at install time because there is no + * pipeline yet when this function runs — the engine gets one per `reload()`, + * and hands it to us as an argument. So the guard lives at the first place a + * pipeline is ever seen. + * + * Failing here means an upgrade moved something and multi-step decoding is + * gone. Stock decoding still produces correct tokens, so the danger is not a + * crash but silence: ~18.4 -> ~9.7 tok/s with nothing in the log to explain + * it. Hence one loud report, and a `fallbacks` count the worker can surface. + */ + const canBurst = (pipeline) => { + const known = supported.get(pipeline); + if (known !== undefined) return known; + + const missing = missingPipelineMembers(pipeline); + supported.set(pipeline, missing.length === 0); + if (missing.length > 0) { + config.fallbacks += 1; + console.error( + "[everything-webgpu] multi-step decoding disabled — falling back to stock " + + "single-step decode. Generation stays correct, throughput roughly halves.\n" + + ` The pipeline is missing ${missing.length} of the internals a burst drives:\n` + + missing.map((line) => ` - ${line}`).join("\n") + + "\n This is what a WebLLM upgrade looks like from here. `npm test` " + + "(webllm-contract) says whether the names are gone from the bundle too.", + ); + onFallback?.({ missing }); + } + return missing.length === 0; + }; + + // A round can end with tokens still buffered — a stop token mid-burst, or an + // interrupt that breaks the caller's loop. Those tokens are already in the KV + // cache, so they must come back out before the next round reuses it. + engine.prefill = async (input, pipeline, chatConfig, genConfig) => { + discardLookahead(pipeline, stateFor(pipeline)); + return basePrefill(input, pipeline, chatConfig, genConfig); + }; + + engine.decode = async (pipeline, genConfig) => { + const state = stateFor(pipeline); + + if (state.queue.length === 0) { + // Before the first burst on this pipeline, not before every one: the + // verdict is cached, so a healthy pipeline pays one property scan for the + // whole conversation. + if (!canBurst(pipeline)) return baseDecode(pipeline, genConfig); + + const burst = burstSize(pipeline, genConfig, config.steps); + if (burst <= 1) return baseDecode(pipeline, genConfig); + + const probe = {}; + const tstart = performance.now(); + state.queue = await sampleBurst(pipeline, genConfig, burst, probe); + const ms = performance.now() - tstart; + + // One burst is one wall-clock cost; its tokens are counted as they drain, + // so a rewound overshoot never inflates the reported rate. + pipeline.decodingTotalTime += ms / 1e3; + pipeline.curRoundDecodingTotalTime += ms / 1e3; + onBurst?.({ steps: burst, tokens: state.queue.length, ms, ...probe }); + } + + const token = state.queue.shift(); + pipeline.decodingTotalTokens += 1; + pipeline.curRoundDecodingTotalTokens += 1; + pipeline.processNextToken(token, genConfig); + + // The burst ran past a stop token; nothing after it was ever emitted. + if (pipeline.stopped() && state.queue.length > 0) discardLookahead(pipeline, state); + }; + + return { + setSteps: (n) => void (config.steps = clampSteps(n)), + get steps() { + return config.steps; + }, + /** Pipelines that failed the contract. Non-zero means the fast path is off. */ + get fallbacks() { + return config.fallbacks; + }, + }; +} + +// ------------------------------------------------------------- burst size --- + +/** + * How many steps may run before the next stop condition *has* to be checked. + * + * `max_tokens` and the context window are countable, so they are clamped rather + * than overshot — which leaves stop tokens as the only reason a burst is ever + * rewound. Returns 1 when multi-step cannot be used at all, routing the caller + * to stock single-step decoding. + */ +export function burstSize(pipeline, genConfig, steps) { + if (steps <= 1) return 1; + + // Per-token CPU feedback: the next step's logits depend on a JS-side decision + // about this step's token, so there is nothing to overlap. + const format = genConfig?.response_format?.type; + if (format === "json_object" || format === "grammar" || format === "structural_tag") return 1; + if (genConfig?.logprobs) return 1; + if (pipeline.logitProcessor !== undefined) return 1; + + const maxTokens = genConfig?.max_tokens; + const untilMax = maxTokens ? maxTokens - pipeline.outputIds.length : Infinity; + const untilContextEnd = + pipeline.slidingWindowSize === -1 + ? pipeline.contextWindowSize - pipeline.filledKVCacheLength + : Infinity; + + return Math.max(1, Math.min(steps, untilMax, untilContextEnd)); +} + +// ----------------------------------------------------------------- burst ---- + +/** + * Runs `steps` forward+sample steps with no GPU->CPU sync between them, then + * pays exactly one. + * + * @returns {Promise} the sampled token ids, in order + */ +async function sampleBurst(pipeline, genConfig, steps, out) { + const { tvm, device } = pipeline; + let probe = null; + const vocab = pipeline.fullVocabSize; + const sampling = resolveSampling(pipeline, genConfig); + + tvm.beginScope(); + let temperatures; + let bias; + let penalty; + /** The last committed token, which seeds step 0. Owned here, not by a scope. */ + let seedTokens; + try { + temperatures = tvm.detachFromCurrentScope( + tvm.empty([1], "float32", device).copyFrom([Math.max(1e-6, sampling.temperature)]), + ); + bias = makeLogitBias(pipeline, sampling); + penalty = makePenalty(pipeline, sampling); + // top_p lives in a tensor the pipeline owns and reuses, set up exactly as + // `sampleTokenFromLogits` does. It is constant for the whole burst. + const topPHost = new Float32Array(pipeline.topPDevice.shape[0]).fill(-1); + const topP = Math.max(sampling.top_p, 1e-5); + pipeline.sampleIndices.forEach((row) => (topPHost[row] = topP)); + pipeline.topPDevice.copyFrom(topPHost); + seedTokens = tvm.detachFromCurrentScope( + tvm.empty([1], "int32", device).copyFrom([pipeline.outputIds[pipeline.outputIds.length - 1]]), + ); + } finally { + tvm.endScope(); + } + let tokens = seedTokens; + + /** + * Sampled ids stay on the device for the whole loop; the host copies happen + * after it, never interleaved with compute. + * + * The order is load-bearing. `flushCommands()` nulls tvmjs's + * `pendingGPUToCPUCopy` whenever it submits an encoder, and every GPU->CPU + * copy calls it. Interleaving copies with compute therefore made each step + * discard the previous step's pending readback, leaving `device.sync()` + * awaiting only the last one — correct in practice only because the + * `mapAsync` promises happen to resolve in FIFO order. Doing all the copies + * after the loop means the first flushes and starts the chain while the rest + * find no pending encoder, so the chain accumulates intact. + */ + const sampledIds = []; + /** One CPU int32[1] per step. All of their reads land on the same poll tick. */ + const staged = []; + + // The decisive probe. The K-step loop below contains no `await`, so it is one + // synchronous JS turn: everything it costs is content-process CPU — command + // encoding, `createBindGroup`, IPC to the GPU process. The `await` after it is + // everything else: GPU execution plus the wait for the next 100 ms poll tick. + // Splitting the two says which one the budget actually goes to. + const gpuCtx = tvm.lib?.webGPUContext; + const dispatchesBefore = gpuCtx?.shaderSubmitCounter ?? 0; + const flushesBefore = countFlushes(gpuCtx); + let forwardDispatches = 0; + const tEncodeStart = performance.now(); + + try { + for (let step = 0; step < steps; step++) { + tvm.beginScope(); + const stepStart = gpuCtx?.shaderSubmitCounter ?? 0; + try { + // `tokens` is owned by `sampledIds` (or is the seed), not by this scope. + const embeddings = pipeline.embed(tokens, pipeline.params); + const batched = embeddings.view([1].concat(embeddings.shape)); + + const states = pipeline.getActiveKVStates(); + const seqIds = tvm.makeShapeTuple([0]); + const inputLen = tvm.makeShapeTuple([1]); + for (const state of states) pipeline.fKVCacheBeginForward(state, seqIds, inputLen); + const forwarded = pipeline.invokeDecode(batched); + for (let i = states.length - 1; i >= 0; i--) pipeline.fKVCacheEndForward(states[i]); + pipeline.filledKVCacheLength += 1; + + // Split the launch count at the forward/sample boundary. The sampling + // tail is `argsort_probs` over the full vocab (248k here), which is a + // multi-pass sort and belongs to the runtime, not the model — so it is + // worth knowing how much of the per-token kernel budget it owns. + forwardDispatches += (gpuCtx?.shaderSubmitCounter ?? 0) - stepStart; + + const logits = forwarded.get(0); + if (bias) { + pipeline.fapplyLogitBias(logits.view([1, vocab]), bias.pos2seqIds, bias.tokenIds, bias.values); + } + if (penalty) { + pipeline.fapplyPenalty( + logits.view([1, vocab]), + penalty.seqIds, + penalty.pos2seqIds, + penalty.tokenIds, + penalty.counts, + penalty.penalties, + ); + } + + const probs = pipeline + .fsoftmaxWithTemperature(logits.view([1, 1, vocab]), temperatures) + .view([1, vocab]); + const sorted = pipeline.fargsortProbs(probs); + const sampled = pipeline.fsampleWithTopP( + sorted.get(0), + sorted.get(1), + tvm.uniform([1], 0, 1, device), + pipeline.sampleIndicesDevice, + pipeline.topPDevice, + ); + + tokens = tvm.detachFromCurrentScope(sampled); + sampledIds.push(tokens); + } finally { + tvm.endScope(); + } + } + + // Every readback together, after all compute: one flush, one intact chain. + tvm.beginScope(); + try { + for (const id of sampledIds) { + staged.push(tvm.detachFromCurrentScope(tvm.empty([1], "int32", tvm.cpu()).copyFrom(id))); + } + } finally { + tvm.endScope(); + } + + // Encoding the copies is still CPU work, so the boundary sits after them. + const tEncoded = performance.now(); + + // The one sync the whole burst pays for. + await device.sync(); + + probe = { + encodeMs: tEncoded - tEncodeStart, + syncMs: performance.now() - tEncoded, + dispatches: (gpuCtx?.shaderSubmitCounter ?? 0) - dispatchesBefore, + forwardDispatches, + flushes: countFlushes(gpuCtx) - flushesBefore, + }; + return staged.map((host) => host.toArray()[0]); + } finally { + if (probe) Object.assign(out, probe); + for (const host of staged) host.dispose(); + for (const id of sampledIds) id.dispose(); + seedTokens?.dispose(); + temperatures.dispose(); + disposeAll(bias); + disposeAll(penalty); + } +} + +/** + * Kernel launches per `flushCommands()`, which decides whether batching tvmjs's + * per-kernel compute passes into one pass is worth anything. + * + * `flushCommands()` submits the pending encoder — so it would also close a + * shared pass — and it fires from `deviceFreeDataSpace`, the buffer copies and + * `sync`. If TVM frees an intermediate between every op then flushes ≈ kernels, + * the pass stream is already chopped up, and there is nothing to merge. tvmjs + * keeps no counter of its own, so wrap the method once per context. + */ +function countFlushes(gpuCtx) { + if (!gpuCtx) return 0; + if (gpuCtx.__ewgpuFlushCount === undefined) { + const base = gpuCtx.flushCommands.bind(gpuCtx); + gpuCtx.__ewgpuFlushCount = 0; + gpuCtx.flushCommands = () => { + gpuCtx.__ewgpuFlushCount += 1; + base(); + }; + } + return gpuCtx.__ewgpuFlushCount; +} + +// ---------------------------------------------------------------- rewind ---- + +/** + * Drops un-emitted lookahead and takes it back out of the KV cache. + * + * `kv_state_popn` is the clean path. If the runtime has not registered it the + * cache cannot be trimmed, so it is thrown away instead: the next round pays a + * full re-prefill (one sync, not one per token) rather than attending over + * tokens the caller never saw. + */ +function discardLookahead(pipeline, state) { + const n = state.queue.length; + state.queue = []; + if (n === 0) return; + + const { tvm } = pipeline; + try { + const popn = getPopN(pipeline); + if (popn) { + tvm.beginScope(); + try { + for (const kvState of pipeline.getActiveKVStates()) { + popn(kvState, tvm.scalar(0, "int64"), tvm.scalar(n, "int32")); + } + } finally { + tvm.endScope(); + } + pipeline.filledKVCacheLength -= n; + return; + } + } catch { + // Fall through: a trim that threw is handled the same as no trim at all. + } + pipeline.resetChat(/* keepStats= */ true); +} + +const popNCache = new WeakMap(); + +function getPopN(pipeline) { + if (popNCache.has(pipeline)) return popNCache.get(pipeline); + let popn = null; + const { tvm } = pipeline; + tvm.beginScope(); + try { + popn = tvm.detachFromCurrentScope(tvm.getGlobalFunc("vm.builtin.kv_state_popn")); + } catch { + popn = null; + } finally { + tvm.endScope(); + } + popNCache.set(pipeline, popn); + return popn; +} + +// ------------------------------------------------------- sampling inputs ---- + +/** + * The subset of `sampleTokenFromLogits`'s config resolution a burst can honour, + * in the same precedence order: the request overrides `mlc-chat-config.json`. + */ +function resolveSampling(pipeline, genConfig) { + const has = (v) => v !== undefined && v !== null; + const pick = (key, fallback) => (has(genConfig?.[key]) ? genConfig[key] : fallback); + return { + temperature: pick("temperature", pipeline.config.temperature), + top_p: pick("top_p", pipeline.config.top_p) ?? 1, + repetition_penalty: pick("repetition_penalty", pipeline.config.repetition_penalty), + frequency_penalty: pick("frequency_penalty", pipeline.config.frequency_penalty) ?? 0, + presence_penalty: pick("presence_penalty", pipeline.config.presence_penalty) ?? 0, + logit_bias: pick("logit_bias", undefined), + }; +} + +/** Static for the whole request, so it is uploaded once and reused every step. */ +function makeLogitBias(pipeline, { logit_bias }) { + const ids = Object.keys(logit_bias ?? {}); + if (ids.length === 0) return null; + const { tvm, device } = pipeline; + const int32 = (values) => + tvm.detachFromCurrentScope(tvm.empty([values.length], "int32", device).copyFrom(values)); + return { + pos2seqIds: int32(new Int32Array(ids.length)), + tokenIds: int32(Int32Array.from(ids, (id) => parseInt(id, 10))), + values: tvm.detachFromCurrentScope( + tvm.empty([ids.length], "float32", device).copyFrom(Float32Array.from(ids, (id) => logit_bias[id])), + ), + }; +} + +/** + * Frozen token history for the burst. + * + * This is the one place multi-step is not equivalent to single-step: tokens + * sampled *within* a burst are not penalised against each other, because their + * ids are still on the GPU. At K=15 the penalty state is at most 15 tokens + * stale. Anything that cannot tolerate that should run with `decodeSteps: 1`. + */ +function makePenalty(pipeline, { repetition_penalty, frequency_penalty, presence_penalty }) { + const active = frequency_penalty !== 0 || presence_penalty !== 0 || (repetition_penalty ?? 1) !== 1; + if (!active) return null; + + const appeared = [...pipeline.appearedTokensFreq.keys()]; + if (appeared.length === 0) return null; + const freqs = [...pipeline.appearedTokensFreq.values()]; + + const { tvm, device } = pipeline; + const int32 = (values) => + tvm.detachFromCurrentScope(tvm.empty([values.length], "int32", device).copyFrom(values)); + return { + seqIds: int32(new Int32Array(1)), + pos2seqIds: int32(new Int32Array(appeared.length)), + tokenIds: int32(Int32Array.from(appeared)), + counts: int32(Int32Array.from(freqs)), + penalties: tvm.detachFromCurrentScope( + tvm + .empty([1, 3], "float32", device) + .copyFrom(new Float32Array([presence_penalty, frequency_penalty, repetition_penalty ?? 1])), + ), + }; +} + +function disposeAll(inputs) { + if (!inputs) return; + for (const tensor of Object.values(inputs)) tensor.dispose(); +} diff --git a/src/engine/pool.js b/src/engine/pool.js new file mode 100644 index 0000000000000000000000000000000000000000..62b0fd8b183b3b16511c5df4c888bc71bef736ef --- /dev/null +++ b/src/engine/pool.js @@ -0,0 +1,439 @@ +/** + * Engine pool + priority scheduler. + * + * Why this exists: the GPU gives us ~10 completion ticks per second, shared by + * every caller (AI.md, "The 10 tok/s ceiling"). Two consequences drive the + * whole design: + * + * - A single generation can never beat ~10 tok/s. When decode is *sync-bound* + * - waiting on the ~100 ms tick rather than on the GPU - N concurrent ones + * each still get ~10 tok/s, because they all wake on the same tick. That + * held for the 0.8B (1.3-2.0x on two engines) and stops holding once the + * model is big enough to keep the GPU busy: measured 1.06x on the 2B, where + * two streams simply run at half speed each. Fan-out is worth having, but + * the second engine earns its memory through isolation, not throughput. + * - Every engine in the pool holds its own copy of the weights, so the pool is + * small and its slots are the scarce resource worth scheduling. + * + * The scheduler is therefore four mechanisms and no more: + * 1. Priority bands, FIFO within a band. + * 2. Session supersession - a new job with the same `session` cancels the + * previous one. This is the ghost-text primitive: every keystroke replaces + * the last request instead of queueing behind it. + * 3. Opt-in preemption - an `interactive` job with no free slot may interrupt + * a running job that declared `preemptible`. The victim resolves with what + * it produced so far, so it is never requeued and can never starve. + * 4. Demand-driven growth - the pool starts at one engine and earns another + * only when a *different* task is waiting on a busy pool. See `#grow`. + * + * Jobs carry a `task`: the unit that owns an engine. Every item of one batch + * shares it, so "translate this page" is one task however many requests it is. + * Slots are handed out per job, but the last free slot is reserved per task - + * otherwise one batch fills the pool and ghost-text starves behind it. + * + * `createEngine` is injected so the scheduler can be tested without a GPU. + */ +import { JOB_KIND, PRIORITY, PRIORITY_ORDER } from "./constants.js"; + +let nextJobId = 0; + +export class EnginePool { + #maxSize; + #createEngine; + #onStateChange; + #slots = []; + #queues = new Map(PRIORITY_ORDER.map((p) => [p, []])); + #bySession = new Map(); + #loading = null; + #growing = null; + /** Set to the failure reason once a grow attempt fails; growth then stops. */ + #growthBlocked = null; + /** + * Bumped by `unload()`. An engine that finishes loading against a stale + * generation is torn down rather than installed — see `load()`. + */ + #generation = 0; + + constructor({ size = 2, createEngine, onStateChange = () => {} }) { + this.#maxSize = Math.max(1, Math.min(4, size)); + this.#createEngine = createEngine; + this.#onStateChange = onStateChange; + } + + /** Engines that exist right now, which is not the same as the cap. */ + get size() { + return this.#slots.length; + } + + get maxSize() { + return this.#maxSize; + } + + get loaded() { + return this.#slots.length > 0; + } + + /** + * Brings up the first engine, and only the first. + * + * The pool used to build every engine here. It no longer does, because an + * engine is a full copy of the weights - measured ~1.6 GB steady state for + * the 0.8B, ~2.4 GB for a 2B - and staging a second one costs that much host + * memory before the GPU ever sees it. On a 16 GB machine that is the + * difference between working and swapping, and it was being paid up front + * whether or not two tasks ever ran at once. `#grow` earns the rest. + */ + async load(onProgress = () => {}) { + if (this.#loading) return this.#loading; + + const generation = this.#generation; + this.#loading = (async () => { + const engine = await this.#createEngine(0, (report) => + onProgress({ ...report, engine: 1, engines: 1 }), + ); + // `unload()` may have run while this was still building — a cancelled + // load, or a model switch. Installing it now would resurrect an engine + // nobody holds a reference to, leaking its worker and a full copy of the + // weights. `#grow()` has always guarded this; `load()` did not. + if (generation !== this.#generation) { + await engine.unload?.().catch(() => {}); + return 0; + } + this.#slots = [{ engine, job: null }]; + return this.#slots.length; + })().finally(() => { + this.#loading = null; + }); + + return this.#loading; + } + + async unload() { + this.#generation += 1; + for (const job of this.#allJobs()) this.#finish(job, { cancelled: true }); + for (const p of this.#queues.values()) p.length = 0; + this.#bySession.clear(); + // A smaller model may well fit where this one did not. + this.#growthBlocked = null; + const slots = this.#slots; + this.#slots = []; + await Promise.all(slots.map((s) => s.engine.unload?.().catch(() => {}))); + } + + /** + * @param {object} spec + * @param {object} spec.params passed straight to `engine.chat.completions.create` + * @param {string} [spec.task] the unit that owns an engine; a whole batch shares one + * @param {string} [spec.session] later jobs with this session supersede earlier ones + * @param {string} [spec.priority] one of PRIORITY + * @param {boolean} [spec.preemptible] may be interrupted by an interactive job + * @param {(chunk: object) => void} [spec.onChunk] WebLLM's chunk, verbatim + * @returns {Promise<{text: string, usage?: object, cancelled?: boolean, preempted?: boolean}>} + */ + submit(spec) { + const priority = PRIORITY_ORDER.includes(spec.priority) ? spec.priority : PRIORITY.NORMAL; + const id = spec.id ?? `job-${++nextJobId}`; + const job = { + id, + // Unlabelled work is its own task, so two bare `chat` calls still compete + // for separate engines the way two different callers would. + task: spec.task ?? spec.session ?? id, + session: spec.session, + priority, + preemptible: Boolean(spec.preemptible), + // What the slot's engine should be asked to do. Everything else about a + // job — priority, supersession, preemption, one-task-one-engine — is the + // same either way, which is why this is a field and not a second pool. + kind: spec.kind === JOB_KIND.EMBEDDING ? JOB_KIND.EMBEDDING : JOB_KIND.CHAT, + params: spec.params, + onChunk: spec.onChunk ?? (() => {}), + text: "", + slot: null, + done: false, + }; + job.promise = new Promise((resolve) => (job.resolve = resolve)); + + if (job.session) { + const previous = this.#bySession.get(job.session); + // Superseded, not queued behind: the keystroke that produced the old + // request is already stale. + if (previous) this.cancel(previous.id); + this.#bySession.set(job.session, job); + } + + this.#queues.get(priority).push(job); + this.#pump(); + this.#emit(); + return job.promise; + } + + /** Cancels by job id or by session key. Returns how many jobs it stopped. */ + cancel(idOrSession) { + let stopped = 0; + for (const job of this.#allJobs()) { + if (job.id !== idOrSession && job.session !== idOrSession) continue; + stopped += 1; + if (job.slot === null) { + this.#dequeue(job); + this.#finish(job, { cancelled: true }); + } else { + job.cancelling = true; + this.#slots[job.slot].engine.interruptGenerate(); + } + } + if (stopped) this.#emit(); + return stopped; + } + + /** + * Pushes a runtime setting to every engine that accepts one. + * + * Separate from `load()` because the settings this carries — `decodeSteps` so + * far — are per-generation knobs, not per-model ones: retuning them must not + * cost a reload of the weights. + */ + configure(patch) { + let applied = 0; + for (const slot of this.#slots) { + if (typeof slot.engine.configure !== "function") continue; + slot.engine.configure(patch); + applied += 1; + } + return applied; + } + + status() { + return { + size: this.#slots.length, + maxSize: this.#maxSize, + growing: this.#growing !== null, + // Non-null once a second engine failed to come up; the UI can say why the + // pool is smaller than the cap instead of looking stuck. + growthBlocked: this.#growthBlocked, + busy: this.#slots.filter((s) => s.job !== null).length, + queued: PRIORITY_ORDER.reduce((n, p) => n + this.#queues.get(p).length, 0), + queuedByPriority: Object.fromEntries( + PRIORITY_ORDER.map((p) => [p, this.#queues.get(p).length]), + ), + }; + } + + // ------------------------------------------------------------ internals --- + + /** + * Adds one engine, but only when a second one would actually buy something. + * + * There is no budget to check against: Firefox implements neither + * `navigator.deviceMemory` nor `performance.memory`, and `storage.estimate()` + * measures disk quota, not RAM. Nothing reports free memory to an extension. + * + * So the pool does not predict, it probes - and it only probes when the + * answer matters. A failed load is taken as the answer and is not retried: + * the retry would cost another full staging pass to learn the same thing. + */ + #grow() { + if (this.#growing || this.#growthBlocked || this.#loading) return; + if (this.#slots.length >= this.#maxSize) return; + if (!this.#crossTaskDemand()) return; + + const index = this.#slots.length; + this.#growing = (async () => { + try { + const engine = await this.#createEngine(index, () => {}); + // `unload()` may have emptied the pool while this was still loading; + // pushing then would resurrect a slot for a model nobody asked for. + if (this.#slots.length === index) this.#slots.push({ engine, job: null }); + else await engine.unload?.().catch(() => {}); + } catch (err) { + // Almost always memory. Staying at the current size is the right + // outcome, not an error owed to whichever job happened to trigger it. + this.#growthBlocked = err?.message ?? String(err); + } + })().finally(() => { + this.#growing = null; + this.#pump(); + this.#emit(); + }); + } + + /** + * True when a queued job belongs to a task that is not already running. + * + * This is the entire growth policy. A second engine exists so that a + * translation and a ghost-text completion can run at once - not to make one + * batch finish sooner. Four queued items of the same batch keep the pool at + * one engine; one queued completion alongside them grows it. + */ + #crossTaskDemand() { + const running = this.#runningTasks(); + if (running.size === 0) return false; + for (const job of this.#queued()) if (!running.has(job.task)) return true; + return false; + } + + #runningTasks() { + const tasks = new Set(); + for (const slot of this.#slots) if (slot.job) tasks.add(slot.job.task); + return tasks; + } + + *#queued() { + for (const p of PRIORITY_ORDER) yield* this.#queues.get(p); + } + + *#allJobs() { + for (const slot of this.#slots) if (slot.job) yield slot.job; + yield* this.#queued(); + } + + #dequeue(job) { + const queue = this.#queues.get(job.priority); + const i = queue.indexOf(job); + if (i >= 0) queue.splice(i, 1); + } + + #nextJob() { + for (const p of PRIORITY_ORDER) { + const queue = this.#queues.get(p); + if (queue.length) return queue[0]; + } + return null; + } + + #pump() { + for (;;) { + const slot = this.#slots.findIndex((s) => s.job === null); + const job = slot >= 0 ? this.#pick() : this.#nextJob(); + if (!job) break; + + if (slot >= 0) { + this.#dequeue(job); + this.#start(slot, job); + continue; + } + + // No slot. Only an interactive job is allowed to take one by force, and + // only from a job that opted in. + if (job.priority !== PRIORITY.INTERACTIVE) break; + const victim = this.#slots.find((s) => s.job?.preemptible && !s.job.cancelling); + if (!victim) break; + victim.job.preempting = true; + victim.job.cancelling = true; + victim.engine.interruptGenerate(); + break; // the freed slot re-enters #pump when the victim settles + } + this.#grow(); + } + + /** + * Which queued job takes a free slot: the first, by priority, whose task is + * not already running. + * + * **One task holds at most one engine.** Letting a batch spread over the pool + * used to be the point - it was worth 1.3-2.0x on the 0.8B, where decode was + * sync-bound and a second stream filled idle GPU. On the 2B it is worth + * 1.06x: the GPU is busy, so two streams of the same work just run at half + * speed each (measured; `engine scaling:` in the e2e). That buys nothing, and + * it costs the thing a second engine is actually for - a page translation + * would sit on both engines and ghost-text would wait behind it. + * + * So the rule is flat, and two runnable tasks are therefore always running at + * once whenever two engines exist. An engine may idle while one task has work + * queued; that is the ~6% being deliberately given up. + * + * `interactive` is the one exception: that band exists for work a human is + * waiting on keystroke-by-keystroke, and it takes a free engine regardless. + */ + #pick() { + const head = this.#nextJob(); + if (!head || head.priority === PRIORITY.INTERACTIVE) return head; + + const running = this.#runningTasks(); + for (const job of this.#queued()) if (!running.has(job.task)) return job; + return null; + } + + #start(slotIndex, job) { + const slot = this.#slots[slotIndex]; + slot.job = job; + job.slot = slotIndex; + job.startedAt = performance.now(); + job.engineIndex = slotIndex; + this.#emit(); + + (async () => { + try { + if (job.kind === JOB_KIND.EMBEDDING) { + // One shot, no stream: an embedding is a single forward pass and + // WebLLM returns the whole OpenAI-shaped response at once. There is + // no decode loop here, so none of the streaming bookkeeping applies — + // but the queueing, priority and preemption above all still did. + const res = await slot.engine.embeddings.create({ ...job.params }); + job.usage = res.usage; + job.embeddings = res.data; + this.#finish(job, { cancelled: Boolean(job.cancelling && !job.preempting) }); + return; + } + + const stream = await slot.engine.chat.completions.create({ + ...job.params, + stream: true, + stream_options: { include_usage: true }, + }); + for await (const chunk of stream) { + const choice = chunk.choices?.[0]; + job.text += choice?.delta?.content ?? ""; + // WebLLM's own "stop" | "length" | "abort" | "tool_calls". Kept + // because a caller cannot otherwise tell a natural stop from a + // `max_tokens` truncation. + if (choice?.finish_reason) job.finishReason = choice.finish_reason; + // Assigned, not merged: WebLLM parses the whole output message at the + // end and emits tool calls complete in one terminal chunk. It never + // streams the OpenAI-style fragments, so there is nothing to + // accumulate and a merge step would be machinery for a wire shape + // that is never produced. + if (choice?.delta?.tool_calls) job.toolCalls = choice.delta.tool_calls; + if (chunk.usage) job.usage = chunk.usage; + // The chunk goes on verbatim. It is already a compliant OpenAI + // envelope carrying id / created / model / logprobs / + // system_fingerprint; rebuilding one downstream only loses fields. + job.onChunk(chunk); + } + this.#finish(job, { + cancelled: Boolean(job.cancelling && !job.preempting), + preempted: Boolean(job.preempting), + }); + } catch (err) { + this.#finish(job, { error: String(err?.message ?? err) }); + } finally { + slot.job = null; + job.slot = null; + this.#pump(); + this.#emit(); + } + })(); + } + + #finish(job, extra) { + if (job.done) return; + job.done = true; + if (this.#bySession.get(job.session) === job) this.#bySession.delete(job.session); + job.resolve({ + id: job.id, + text: job.text, + usage: job.usage, + ...(job.embeddings ? { embeddings: job.embeddings } : {}), + ...(job.toolCalls ? { toolCalls: job.toolCalls } : {}), + // Interrupted work reports "abort" whatever the stream last said. + finishReason: extra.cancelled || extra.preempted ? "abort" : job.finishReason, + engineIndex: job.engineIndex, + startedAt: job.startedAt, + finishedAt: performance.now(), + ...extra, + }); + } + + #emit() { + this.#onStateChange(this.status()); + } +} diff --git a/src/engine/prefetch.js b/src/engine/prefetch.js new file mode 100644 index 0000000000000000000000000000000000000000..794b16c8f4e23a2705698b40f0086234fda335c2 --- /dev/null +++ b/src/engine/prefetch.js @@ -0,0 +1,200 @@ +/** + * `prefetch(modelId)` — fill the cache without building an engine. + * + * The gap this closes: WebLLM only ever downloads a model as a side effect of + * `reload()`, which instantiates the wasm and needs a GPU before it will fetch + * a single weight shard. So "warm the cache during onboarding, decide about the + * GPU later" is not expressible — and it is the thing an app wants to do while + * the user is still reading the welcome screen. + * + * We already write these exact caches for injected models (`ingest.js`); this is + * the same write plan with the bytes arriving over the network instead of off + * disk. + * + * ## The one dangerous part, and how it is closed + * + * To fetch the artifacts ourselves we must know their URLs, which means + * applying HuggingFace's `/resolve/main/` rule — the same rule + * [ARCHIVE.md](../../ARCHIVE.md) records *removing* a copy of, and that §2a of + * the roadmap says not to derive at registration time. + * + * Those decisions still hold and this does not contradict them: they are about + * not deriving a URL that WebLLM will derive again at load, which double-applies + * it. Here WebLLM is not in the loop at all — we are the loader — so there is + * nothing to double-apply. + * + * What makes it *safe* is that we do not trust our own derivation. A key that is + * off by one character writes a cache WebLLM's loader will never look in, and + * the symptom is the worst kind: prefetch reports success and the user downloads + * the model twice. So every prefetch ends by asking **WebLLM's own + * `hasModelInCache`** — which derives the URL through the very function we are + * mirroring — whether the model is really there. If it says no, this throws + * instead of claiming success. + * + * `webllm-contract.test.mjs` additionally pins the rule against the bundle, so + * an upstream change to the URL scheme fails a test rather than a user's + * download. + */ +import { ERROR, EngineError } from "./errors.js"; +import { + CHAT_CONFIG, + CONTENT_TYPES, + LEGACY_TENSOR_MANIFEST, + TENSOR_MANIFEST, +} from "./ingest.js"; +import { CACHE_CONFIG, CACHE_MODEL, CACHE_WASM, isInjected } from "./model-store.js"; + +/** + * WebLLM's `cleanModelUrl`, mirrored deliberately. + * + * Kept character-for-character with the bundle's version (see the contract + * test) because the whole point is to produce the same cache keys its loader + * will look for. + */ +export function resolveModelUrl(modelUrl) { + let url = modelUrl + (modelUrl.endsWith("/") ? "" : "/"); + if (!url.match(/.+\/resolve\/.+\//)) url += "resolve/main/"; + return new URL(url).href; +} + +/** + * @param {object} opts + * @param {string} opts.modelId + * @param {object} opts.record the merged app-config entry: `model`, `model_lib` + * @param {(p: {phase: string, done: number, total: number, label: string}) => void} [opts.onProgress] + * @param {AbortSignal} [opts.signal] + * @param {(url: string, init?: object) => Promise} [opts.fetchImpl] + * @returns {Promise<{modelId: string, files: number, bytes: number, alreadyCached: boolean}>} + */ +export async function prefetchModel({ + modelId, + record, + onProgress = () => {}, + signal, + fetchImpl = globalThis.fetch, +}) { + if (!record?.model_lib) { + throw new EngineError( + ERROR.BAD_REQUEST, + `"${modelId}" has no \`model_lib\`, so there is nothing to prefetch from. ` + + "A remote source needs one; see load(url, { modelLib }).", + { modelId }, + ); + } + + const base = resolveModelUrl(record.model); + const get = async (url, what) => { + if (signal?.aborted) throw aborted(modelId); + const res = await fetchImpl(url, signal ? { signal } : undefined).catch((err) => { + throw new EngineError(ERROR.GENERATION_FAILED, `Prefetch could not reach ${what}: ${err?.message ?? err}`, { + modelId, + url, + }); + }); + if (!res.ok) { + throw new EngineError( + ERROR.UNKNOWN_MODEL, + `Prefetch got ${res.status} for ${what} at ${url}. ` + + "Check the model's base URL — a 404 here usually means the id or the URL is wrong.", + { modelId, url, status: res.status }, + ); + } + return res; + }; + + // The config first: it names the tokenizer files, so the plan cannot be built + // without it. Same order `reload()` uses. + onProgress({ phase: "manifest", done: 0, total: 1, label: CHAT_CONFIG }); + const configRes = await get(base + CHAT_CONFIG, CHAT_CONFIG); + const configBytes = await configRes.arrayBuffer(); + const chatConfig = parseJson(configBytes, CHAT_CONFIG, modelId); + + // `tensor-cache.json`, falling back to the legacy name, exactly as ingest does. + let manifestName = TENSOR_MANIFEST; + let manifestRes = await fetchImpl(base + TENSOR_MANIFEST, signal ? { signal } : undefined).catch(() => null); + if (!manifestRes?.ok) { + manifestName = LEGACY_TENSOR_MANIFEST; + manifestRes = await get(base + LEGACY_TENSOR_MANIFEST, "the weight index"); + } + const manifestBytes = await manifestRes.arrayBuffer(); + const manifest = parseJson(manifestBytes, manifestName, modelId); + + const shards = (manifest.records ?? []).map((r) => r.dataPath).filter(Boolean); + if (shards.length === 0) { + throw new EngineError( + ERROR.UNKNOWN_MODEL, + `${manifestName} at ${base} lists no weight shards, so this is not an MLC model directory.`, + { modelId, url: base + manifestName }, + ); + } + + const tokenizers = (Array.isArray(chatConfig.tokenizer_files) ? chatConfig.tokenizer_files : []).filter( + (n) => n === "tokenizer.json" || n === "tokenizer.model", + ); + + const plan = [ + { scope: CACHE_CONFIG, url: base + CHAT_CONFIG, body: configBytes, type: CONTENT_TYPES.json }, + { scope: CACHE_MODEL, url: base + manifestName, body: manifestBytes, type: CONTENT_TYPES.json }, + ...tokenizers.map((name) => ({ + scope: CACHE_MODEL, + url: base + name, + type: name.endsWith(".json") ? CONTENT_TYPES.json : CONTENT_TYPES.bin, + })), + ...shards.map((p) => ({ scope: CACHE_MODEL, url: new URL(p, base).href, type: CONTENT_TYPES.bin })), + // Verbatim, never derived — `model_lib` is a literal URL on the record and + // is not even on the same origin as the weights for any prebuilt model. + { scope: CACHE_WASM, url: record.model_lib, type: CONTENT_TYPES.wasm }, + ]; + + const openCaches = new Map(); + const cacheFor = async (scope) => { + if (!openCaches.has(scope)) openCaches.set(scope, await caches.open(scope)); + return openCaches.get(scope); + }; + + let bytes = 0; + let done = 0; + for (const item of plan) { + if (signal?.aborted) throw aborted(modelId); + const cache = await cacheFor(item.scope); + onProgress({ phase: "downloading", done, total: plan.length, label: basename(item.url) }); + + // Skip what is already there: a resumed prefetch should cost only the + // remainder, the same way a resumed `load()` does. + if (item.body === undefined && (await cache.match(new Request(item.url)))) { + done += 1; + continue; + } + + const body = item.body ?? (await (await get(item.url, basename(item.url))).arrayBuffer()); + bytes += body.byteLength; + await cache.put( + new Request(item.url), + new Response(body, { status: 200, headers: { "Content-Type": item.type } }), + ); + done += 1; + } + onProgress({ phase: "downloading", done, total: plan.length, label: "done" }); + + return { modelId, files: plan.length, bytes, alreadyCached: false }; +} + +const aborted = (modelId) => + new EngineError(ERROR.ABORTED, `Prefetch of "${modelId}" was aborted.`, { modelId }); + +function parseJson(buffer, what, modelId) { + try { + return JSON.parse(new TextDecoder().decode(buffer)); + } catch { + throw new EngineError( + ERROR.UNKNOWN_MODEL, + `${what} for "${modelId}" is not valid JSON — the URL is probably not an MLC model directory.`, + { modelId, what }, + ); + } +} + +const basename = (url) => url.split("/").pop() || url; + +/** Only for injected models: they are in the cache before they are ever registered. */ +export const isAlreadyLocal = (record) => isInjected(record); diff --git a/src/engine/recipes.js b/src/engine/recipes.js new file mode 100644 index 0000000000000000000000000000000000000000..09fc0d969ef2effb58c03c54cb33152021827ad9 --- /dev/null +++ b/src/engine/recipes.js @@ -0,0 +1,281 @@ +/** + * The three shapes of work, as one call each. + * + * `complete()` and `chat.completions.create()` can already express all of this. + * What they cannot do is stop a caller getting the *scheduling* wrong, and the + * scheduling is the part that is easy to get wrong and invisible when you do: + * ghost text that lags a keystroke behind, a conversation whose turns fight + * each other for engines, a one-shot that superseded the last one because it + * reused a session key. Those are the bugs AI.md's "Getting these wrong" table + * is made of, and every row of it is a scheduling mistake rather than a + * generation one. + * + * So these are not wrappers that save typing. Each one is a *policy*: + * + * ask() one-shot, its own task, no session, nothing to supersede + * conversation() one stable task for every turn, history the caller can see + * ghostText() debounce + one session key + interactive + drop-if-stale + * + * ## What they deliberately do not do + * + * **They author no prompts.** `ask()` and `conversation()` carry the caller's + * own text through as a message; `ghostText()` *requires* a `prompt` function + * and has no default for it. This is AI.md's load-bearing rule, and the reason + * is not purity: prompts are model-specific, and switching this project's own + * build from `Qwen3.5-0.8B` to `Qwen3.8-2B-Distill` changed the conversation + * template and made every reply open with a `` block. A prompt that + * lives in the caller survives that. One baked in here would have to be + * rewritten and re-shipped to every caller. + */ +import { ERROR, EngineError } from "./errors.js"; +import { PRIORITY } from "./constants.js"; + +/** + * A single question, with nothing kept afterwards. + * + * Its own task and no session, so two `ask()`s in flight never supersede each + * other and never queue behind one another for the same engine — which is what + * would happen if this shared a session key with anything else. + * + * @param {import("./engine.js").ScheduledEngine} engine + * @param {string | Array<{role: string, content: string}>} input + * @param {object} [opts] anything `complete()` takes; `onDelta` to stream + * @returns {Promise} the reply text + */ +export async function ask(engine, input, { onDelta, ...opts } = {}) { + const messages = toMessages(input, "ask"); + const { text } = await engine.complete( + { priority: PRIORITY.NORMAL, ...opts, messages }, + onDelta, + ); + return text; +} + +/** + * A multi-turn conversation that remembers its own history. + * + * Two things this gets right that hand-rolled history usually does not: + * + * **One task for the whole conversation.** Every turn carries the same `task`, + * so a conversation holds at most one engine and a long reply can never occupy + * the pool while another conversation waits. Turns within it are serialised, + * which is what a conversation means anyway. + * + * **A bounded history, by default.** There is no cross-turn KV reuse on this + * stack (AI.md: every turn re-prefills the whole history at ~5.27 ms/token), so + * an unbounded conversation gets quadratically slower and a turn near the 4096 + * limit waits ~22 s for its first token. `keep` bounds it. Set `keep: Infinity` + * to opt out, having read that sentence. + * + * @param {import("./engine.js").ScheduledEngine} engine + * @param {object} [opts] + * @param {string} [opts.system] a system message, prepended and never trimmed + * @param {number} [opts.keep] how many *exchanges* of history to carry + * @param {string} [opts.task] defaults to a fresh id + */ +export function conversation(engine, { system, keep = 12, task, ...defaults } = {}) { + if (typeof keep !== "number" || (keep < 1 && keep !== Infinity)) { + throw new EngineError( + ERROR.BAD_REQUEST, + `conversation({ keep }) must be a positive number of exchanges, or Infinity. Got ${JSON.stringify(keep)}.`, + { keep }, + ); + } + + const id = task ?? `conversation-${++counter}`; + /** User/assistant messages only; `system` is held apart so trimming cannot eat it. */ + let turns = []; + /** Turns are serialised: overlapping ones would interleave history. */ + let inFlight = Promise.resolve(); + + async function turn(content, onDelta) { + const sent = [...api.messages, { role: "user", content }]; + const { text, finishReason } = await engine.complete( + { ...defaults, task: id, priority: defaults.priority ?? PRIORITY.NORMAL, messages: sent }, + onDelta, + ); + // Appended only on success: a failed turn must not leave the history + // holding a question the model never answered. + turns.push({ role: "user", content }, { role: "assistant", content: text }); + trim(); + return { text, finishReason }; + } + + const api = { + /** The messages as they would be sent, including the system message. */ + get messages() { + return system ? [{ role: "system", content: system }, ...turns] : [...turns]; + }, + + /** Exchanges currently retained. */ + get length() { + return Math.ceil(turns.length / 2); + }, + + /** + * Say something and get the reply, with the exchange appended to history. + * + * Serialised against the previous turn: a conversation whose turns + * overlapped would interleave history and produce replies to the wrong + * question. + * + * @param {string} content + * @param {(delta: string) => void} [onDelta] + */ + async say(content, onDelta) { + if (typeof content !== "string" || content.length === 0) { + throw new EngineError(ERROR.BAD_REQUEST, "conversation.say() needs a non-empty string.", { + received: typeof content, + }); + } + const run = inFlight.then(() => turn(content, onDelta)); + // The chain must survive a failed turn: `.catch` here keeps the *queue* + // moving without swallowing the rejection the caller is awaiting. + inFlight = run.catch(() => {}); + return run; + }, + + /** Forget the history. The system message and settings survive. */ + reset() { + turns = []; + return api; + }, + + /** Drop the history in, e.g. when restoring a saved conversation. */ + restore(messages) { + turns = messages.filter((m) => m.role !== "system").map((m) => ({ ...m })); + trim(); + return api; + }, + }; + + function trim() { + if (keep === Infinity) return; + const max = keep * 2; + if (turns.length > max) turns = turns.slice(turns.length - max); + } + + return api; +} + +/** + * Ghost text: the scheduling discipline, with the prompt left to the caller. + * + * Every part of this exists because of a specific way ghost text goes wrong: + * + * | | | + * | --- | --- | + * | one stable `session` | a fresh id per keystroke makes every stale request still run — AI.md's first "getting these wrong" row | + * | `interactive` priority | it is the one band that may preempt work that opted in | + * | debounce | a request per keystroke queues faster than the GPU drains | + * | short `max_tokens` | ghost text is a few words; paying for more is pure latency | + * | resolves `null` when stale | so a superseded suggestion **cannot** be rendered by mistake | + * + * That last one is the difference between this and a wrapper. The engine + * already supersedes stale requests; what a caller still has to remember is not + * to paint the answer that comes back. Returning `null` removes the choice. + * + * @param {import("./engine.js").ScheduledEngine} engine + * @param {object} opts + * @param {(context: any) => string | Array} opts.prompt **required** — + * builds the messages. Never defaulted: see the module header. + * @param {number} [opts.debounceMs] + * @param {number} [opts.maxTokens] + * @param {string} [opts.session] + */ +export function ghostText(engine, { prompt, debounceMs = 120, maxTokens = 24, session, ...defaults } = {}) { + if (typeof prompt !== "function") { + throw new EngineError( + ERROR.BAD_REQUEST, + "ghostText({ prompt }) needs a function that turns your editor context into messages. " + + "It has no default on purpose: prompts are model-specific and belong to whoever owns the feature.", + { received: typeof prompt }, + ); + } + + const key = session ?? `ghost-${++counter}`; + let generation = 0; + let timer = null; + /** + * Settles the debounce wait of the call currently holding the timer. + * + * Without this, cancelling that timer left its `await` with nothing to + * resolve it — every superseded keystroke leaked a promise that never + * settled, and `Promise.all` over a burst of them hung forever. A superseded + * waiter has to be *woken and told it lost*, not merely disarmed. + */ + let wake = null; + + const stopWaiting = (quiet) => { + clearTimeout(timer); + timer = null; + const settle = wake; + wake = null; + settle?.(quiet); + }; + + const api = { + /** + * Ask for a suggestion. Debounced, superseding, and `null` when stale. + * + * @param {any} context whatever `prompt` takes + * @returns {Promise} `null` if superseded or cancelled + */ + async suggest(context) { + const mine = ++generation; + + if (debounceMs > 0) { + stopWaiting(false); // the previous waiter loses, and is told so + const quiet = await new Promise((resolve) => { + wake = resolve; + timer = setTimeout(() => stopWaiting(true), debounceMs); + }); + // A newer keystroke landed while waiting; that request owns the session. + if (!quiet || mine !== generation) return null; + } + + const messages = toMessages(prompt(context), "ghostText's prompt()"); + const result = await engine.complete({ + max_tokens: maxTokens, + ...defaults, + messages, + session: key, + priority: PRIORITY.INTERACTIVE, + }); + + // Two ways to be stale, and both must return null: the engine superseded + // us (`cancelled`), or a newer suggest() started while we generated. + if (result.cancelled || result.preempted || mine !== generation) return null; + return result.text; + }, + + /** On blur, or on accept. Cancels in flight and invalidates anything pending. */ + cancel() { + generation += 1; + stopWaiting(false); + return engine.cancel(key); + }, + }; + + return api; +} + +let counter = 0; + +/** The one place caller text becomes a message, so no verb invents its own shape. */ +function toMessages(input, who) { + if (typeof input === "string") { + if (input.length === 0) { + throw new EngineError(ERROR.BAD_REQUEST, `${who} was given an empty string.`, { who }); + } + return [{ role: "user", content: input }]; + } + if (Array.isArray(input) && input.length > 0 && input.every((m) => m?.role && m?.content !== undefined)) { + return input; + } + throw new EngineError( + ERROR.BAD_REQUEST, + `${who} needs a string, or a non-empty array of { role, content } messages.`, + { who, received: Array.isArray(input) ? "array" : typeof input }, + ); +} diff --git a/src/engine/sources.js b/src/engine/sources.js new file mode 100644 index 0000000000000000000000000000000000000000..338e39f07da4442d64f6f55afca23ac36b92cef9 --- /dev/null +++ b/src/engine/sources.js @@ -0,0 +1,197 @@ +/** + * What `load()` was handed, decided before anything touches the network. + * + * `load` absorbs `registerModel` and `ingestModelFolder`, so one call has to + * tell four things apart: a model id, a URL to fetch from, an explicit remote + * spec, and a folder off disk. Keeping that decision here — pure, synchronous, + * no store and no engine — means it can be tested exhaustively without a GPU, + * and means `load()` reads as "classify, then act". + * + * Two rules this deliberately does **not** implement, both measured rather than + * assumed: + * + * - **`modelLib` is never guessed.** `-webgpu.wasm` matches 0 of the + * 163 prebuilt models (real names carry a `_cs1k`-style suffix and drop the + * `-MLC`), and 0 of 163 host the lib on the same origin as the weights — + * they live on `raw.githubusercontent.com`. A guess would be wrong every + * time and would surface as a 404 deep inside WebLLM's loader, which is the + * worst place to learn it. So a URL source without `modelLib` fails here, + * immediately, with a sentence saying why. + * - **`/resolve/main/` is never derived for HuggingFace URLs.** WebLLM's + * `cleanModelUrl` already appends it when the URL does not match + * `.+/resolve/.+/`. Deriving it here would re-introduce exactly the + * duplication ARCHIVE.md records removing. The URL is passed through. + */ +import { ERROR, EngineError } from "./errors.js"; + +/** @typedef {"id"|"remote"|"register"|"files"} SourceKind */ + +export const SOURCE_KIND = { + ID: "id", + REMOTE: "remote", + REGISTER: "register", + FILES: "files", +}; + +/** A scheme, as `absolutize()` recognises one. */ +const HAS_SCHEME = /^[a-z][a-z0-9+.-]*:/i; + +/** + * A string is a location rather than an id when it carries a scheme or is + * explicitly path-shaped. No prebuilt model id contains `/` or `:`, so this + * never steals a real id — and it means `load("/models/foo/")` is understood + * as the URL it obviously is instead of being looked up and reported missing. + */ +export const looksLikeUrl = (s) => + HAS_SCHEME.test(s) || s.startsWith("/") || s.startsWith("./") || s.startsWith("../"); + +/** Duck-typed: `DataTransfer` and `FileList` do not exist in Node. */ +export const isDataTransfer = (v) => !!v && typeof v === "object" && "items" in v && "files" in v; +export const isFileList = (v) => + !!v && + typeof v === "object" && + typeof v.length === "number" && + !Array.isArray(v) && + (v.length === 0 || typeof v[0]?.name === "string"); + +/** + * The local label for a model loaded from a URL. + * + * Safe to derive, unlike `modelLib`: an id is a key in our own registry, never + * a path anything fetches, so a wrong guess is visible immediately and costs + * nothing. `.../mlc-ai/Foo-MLC` and `.../models/foo/` both give the last real + * segment. + */ +export function idFromUrl(url) { + const path = HAS_SCHEME.test(url) ? safeUrlPath(url) : url; + const segments = path.split("/").filter((s) => s && s !== "." && s !== ".."); + const last = segments.at(-1); + return last ? decodeURIComponent(last) : null; +} + +function safeUrlPath(url) { + try { + return new URL(url).pathname; + } catch { + return url; + } +} + +/** + * @param {string | object} src + * @param {{id?: string, modelLib?: string}} [opts] + * @returns {{kind: SourceKind, modelId?: string, model?: string, modelLib?: string, files?: unknown}} + */ +export function classifySource(src, opts = {}) { + if (src === null || src === undefined || src === "") { + throw new EngineError( + ERROR.BAD_REQUEST, + "load() needs a model id, a URL, `{ model, modelLib }`, or a folder " + + "(`{ files }`, a FileList, or a drop event's DataTransfer).", + ); + } + + if (typeof src === "string") { + if (!looksLikeUrl(src)) return { kind: SOURCE_KIND.ID, modelId: src }; + return remote(src, opts.modelLib, opts.id ?? idFromUrl(src)); + } + + // Folder routes first: a DataTransfer also has other properties, and `files` + // is the field that decides, exactly as `registerModel` already reads it. + if (isDataTransfer(src) || isFileList(src) || Array.isArray(src)) { + return { kind: SOURCE_KIND.FILES, files: src, modelId: opts.id }; + } + + if (typeof src === "object") { + if (src.files !== undefined) { + if (src.model || src.modelLib) { + throw new EngineError( + ERROR.BAD_REQUEST, + "load() takes either `files` (local, never fetched) or `model`/`modelLib` " + + "(a base URL to fetch), not both.", + ); + } + return { kind: SOURCE_KIND.FILES, files: src.files, modelId: src.modelId ?? opts.id }; + } + if (src.model !== undefined) { + // The id falls back to the URL exactly as the string form's does — the + // two spellings of "a model at this URL" must not disagree about it. + const spec = remote( + src.model, + src.modelLib ?? opts.modelLib, + src.modelId ?? opts.id ?? idFromUrl(String(src.model)), + ); + return { ...spec, kind: SOURCE_KIND.REGISTER }; + } + } + + throw new EngineError( + ERROR.BAD_REQUEST, + "load() did not recognise that source. Pass a model id, a URL, " + + "`{ model, modelLib }`, or a folder (`{ files }`, a FileList, or a DataTransfer).", + { received: typeof src }, + ); +} + +/** The one place a URL source is validated, so both string and object forms agree. */ +function remote(model, modelLib, modelId) { + if (!modelLib) { + throw new EngineError( + ERROR.BAD_REQUEST, + `Loading "${model}" from a URL needs \`modelLib\` — the compiled .wasm — and it cannot be ` + + "guessed: across all 163 prebuilt models, none has a lib name derivable from its id and " + + "none hosts the lib on the same origin as the weights. " + + "Pass load(url, { modelLib: \"https://.../foo-webgpu.wasm\" }).", + { model }, + ); + } + if (!modelId) { + throw new EngineError( + ERROR.BAD_REQUEST, + `Could not derive a model id from "${model}". Pass one as load(src, { id }).`, + { model }, + ); + } + return { kind: SOURCE_KIND.REMOTE, model, modelLib, modelId }; +} + +/** + * Ids close enough to be a typo, so an unknown id can say "did you mean". + * + * Cheap on purpose — case-insensitive substring both ways, then a bounded edit + * distance. The list is at most a few hundred entries and this runs once, on a + * path that is already about to throw. + */ +export function nearMatches(wanted, available, limit = 3) { + const needle = wanted.toLowerCase(); + const scored = available + .map((id) => { + const hay = id.toLowerCase(); + if (hay === needle) return { id, score: 0 }; + if (hay.includes(needle) || needle.includes(hay)) return { id, score: 1 }; + return { id, score: 2 + editDistance(needle, hay) }; + }) + // Beyond this a "suggestion" is noise, and a wrong suggestion is worse than + // none — it sends the reader looking for a model that was never the point. + .filter(({ score, id }) => score <= 2 + Math.ceil(Math.max(id.length, needle.length) / 3)) + .sort((a, b) => a.score - b.score || a.id.localeCompare(b.id)); + + return scored.slice(0, limit).map(({ id }) => id); +} + +/** Levenshtein, two rows. */ +function editDistance(a, b) { + let prev = Array.from({ length: b.length + 1 }, (_, i) => i); + for (let i = 1; i <= a.length; i++) { + const row = [i]; + for (let j = 1; j <= b.length; j++) { + row[j] = Math.min( + prev[j] + 1, + row[j - 1] + 1, + prev[j - 1] + (a[i - 1] === b[j - 1] ? 0 : 1), + ); + } + prev = row; + } + return prev[b.length]; +} diff --git a/src/manager/manager.css b/src/manager/manager.css new file mode 100644 index 0000000000000000000000000000000000000000..d029bf7b6c8cc71f657d676deefe8ce74a870037 --- /dev/null +++ b/src/manager/manager.css @@ -0,0 +1,83 @@ +.wrap { max-width: 820px; margin: 0 auto; padding: 28px 24px 72px; display: flex; flex-direction: column; gap: 30px; } + +header { padding-bottom: 14px; border-bottom: 1px solid var(--line); } +section { display: flex; flex-direction: column; gap: 10px; } + +/* ------------------------------------------------------------ engine API --- */ +/* First section on the page: an integrator opening this page is here for the + ID and a snippet, not to re-read the model list. */ +#api { padding: 18px; border: 1px solid var(--line-strong); border-radius: 12px; background: var(--panel); } + +.idbar { display: flex; gap: 8px; align-items: center; } +.idbar > code { + padding: 7px 12px; + border: 1px solid var(--line); + border-radius: 8px; + background: var(--bg); + overflow-x: auto; + white-space: nowrap; +} + +.tabs { display: flex; gap: 4px; margin-top: 4px; } +.tab { border-color: transparent; color: var(--muted); padding: 5px 11px; border-radius: 99px; } +.tab:hover:not(.active) { background: var(--panel-strong); } +.tab.active { background: var(--accent); border-color: var(--accent); color: var(--on-accent); font-weight: 500; } + +.codewrap { position: relative; } +.codewrap .copy { + position: absolute; + top: 8px; + right: 8px; + padding: 4px 10px; + font-size: 12px; + background: var(--bg); + opacity: 0; + transition: opacity .12s ease; +} +.codewrap:hover .copy, .codewrap .copy:focus-visible { opacity: 1; } + +/* ---------------------------------------------------------------- models --- */ +#drop { + border: 1.5px dashed var(--line-strong); + border-radius: 12px; + padding: 28px 20px; + text-align: center; + display: flex; + flex-direction: column; + gap: 10px; + align-items: center; + transition: border-color .15s, background .15s; +} +#drop.over { border-color: var(--accent); background: var(--panel-strong); } +#drop p { margin: 0; max-width: 560px; } + +table { width: 100%; border-collapse: collapse; } +th, td { text-align: left; padding: 8px; border-bottom: 1px solid var(--line); vertical-align: middle; } +th { color: var(--muted); font-weight: 500; font-size: 11px; text-transform: uppercase; letter-spacing: .05em; } +td.mono { font-size: 12px; } +table.flags td { border: none; padding: 3px 12px 3px 0; } + +.grid { display: grid; grid-template-columns: repeat(auto-fit, minmax(200px, 1fr)); gap: 12px; } +label { display: flex; flex-direction: column; gap: 5px; color: var(--muted); font-size: 12px; } +label.block { width: 100%; } +label input, label textarea { color: var(--ink); background: var(--panel); } +textarea { resize: vertical; } + +pre { + background: var(--bg); + border: 1px solid var(--line); + border-radius: 10px; + padding: 14px; + overflow-x: auto; + margin: 0; + line-height: 1.55; +} +code { font-family: ui-monospace, "SF Mono", Menlo, monospace; } + +.pill { padding: 2px 8px; border-radius: 99px; font-size: 10px; border: 1px solid var(--line-strong); color: var(--muted); } +.pill.ok { color: var(--accent); border-color: var(--accent); } +.pill.bad { color: var(--ink); border-color: var(--ink); } + +/* There is no fourth colour to warn in, so a warning is stated in the accent + and carried by the sentence rather than by hue alone. */ +.warn { color: var(--accent); } diff --git a/src/manager/manager.html b/src/manager/manager.html new file mode 100644 index 0000000000000000000000000000000000000000..7accbf7cf0bf4471efa426ee555dc67622269a65 --- /dev/null +++ b/src/manager/manager.html @@ -0,0 +1,122 @@ + + + + + Everything WebGPU — Models + + + + +
+ +
+ +

Everything WebGPU

+ + + +
+ + + +
+

Engine API

+
+ + +
+

+ Any Firefox extension drives this engine by messaging that ID — runtime.sendMessage for + request/response, runtime.connect for streaming. Send raw requests: the engine schedules a + shared GPU, it does not author prompts, so prompts stay on your side and survive a model change. +

+ +
+ + + +
+

+
+ +
+
+

+ One task holds one engine, so a whole batch occupies one and can never starve a completion. + A second engine is built only when a second task is waiting, and on this model it is worth ~1.06x — it + buys isolation, not speed. Full reference in AI.md. +

+
+ +
+

Add a model

+
+ Drop a compiled MLC model folder here +

+ Files are copied straight into this extension's Cache Storage — nothing is downloaded. + The folder needs mlc-chat-config.json, tensor-cache.json + (or ndarray-cache.json), every params_shard_*.bin, + tokenizer.json, and exactly one *-webgpu.wasm model library. +

+ + +
+ + +
+ +
+

Local models

+ + + + + +
ModelSizeShardsLibraryCache
+ +

Storage in use: —

+
+ +
+

Generation defaults

+
+ + + + +
+

+

+ + +
+ + +
+
+ +
+

Firefox setup (macOS)

+

WebGPU/Metal must be enabled in about:config before a model can load:

+ + + + + + +
dom.webgpu.enabledtrueExposes navigator.gpu.
gfx.webgpu.ignore-blocklisttrueOnly if your Mac's GPU is blocklisted.
dom.webgpu.service-workers.enabledtrueHarmless; needed on builds that gate worker contexts.
+

Restart Firefox after changing these. Nightly on Apple Silicon is the best-tested combination.

+
+ +
+ + + diff --git a/src/manager/manager.js b/src/manager/manager.js new file mode 100644 index 0000000000000000000000000000000000000000..2534a8a5916a5961519d25f84758c04bbda91abc --- /dev/null +++ b/src/manager/manager.js @@ -0,0 +1,350 @@ +/** Manager page: model ingestion, registry maintenance, engine settings, setup help. */ +import { ENGINE_STATE, OP, PORT_NAME, PORT_OP, PRIORITY, PROTOCOL, request } from "../adapters/protocol.js"; +import { webExtensionStorage } from "../adapters/webext.js"; +import { ModelStore, formatBytes } from "../engine/model-store.js"; +import { filesFromDataTransfer, filesFromInput, ingestModelFolder } from "../engine/ingest.js"; +import { clampSteps } from "../engine/multistep.js"; + +/** + * The manager page writes into the same Cache Storage and the same + * `browser.storage.local` the background page reads, so it builds its own + * ModelStore over the same adapter rather than routing registry edits through + * the engine. Only generation goes over the wire. + */ +const store = new ModelStore(webExtensionStorage()); + +const $ = (id) => document.getElementById(id); + +const port = browser.runtime.connect({ name: PORT_NAME }); +port.onMessage.addListener((msg) => { + if (msg?.protocol === PROTOCOL && msg.op === PORT_OP.ENGINE_STATE) renderEngine(msg.state); +}); + +let lastEngineState = {}; + +function renderEngine(state) { + lastEngineState = state; + $("dot").className = `dot ${state.status}`; + // `size` is engines that exist, which trails `maxSize` until a second task + // asks for one — so say both, or a pool of 1 under a cap of 2 reads as a bug. + const pool = state.pool?.size + ? ` · pool ${state.pool.busy}/${state.pool.size} busy, ${state.pool.queued} queued` + + (state.pool.growthBlocked + ? ` · stayed at ${state.pool.size} (${state.pool.growthBlocked})` + : state.pool.size < state.pool.maxSize + ? ` · up to ${state.pool.maxSize} on demand` + : "") + : ""; + const loading = state.status === ENGINE_STATE.LOADING; + // Show WebLLM's own report verbatim while loading — shard counts, MB and + // elapsed seconds are the only feedback there is during a ~48 s load, and a + // bare percentage hides all of it. + $("engineStatus").textContent = + state.error ? state.error : + state.status === ENGINE_STATE.READY ? `loaded: ${state.modelId}${pool}` : + loading ? state.progress?.text ?? `loading ${state.modelId}…` : + "idle"; + + $("engineBar").hidden = !loading; + $("engineProgress").style.width = `${Math.round((state.progress?.progress ?? 0) * 100)}%`; + const hint = $("engineHint"); + if (!loading) { + hint.hidden = true; + } else { + const secs = state.progress?.timeElapsed; + // Not WebLLM's stock "first visit populates the cache" line: these weights + // were injected by drag-and-drop, so nothing is ever downloaded. + hint.textContent = + `${secs ? `${secs}s elapsed · ` : ""}reading from local cache, no network` + + ((state.progress?.progress ?? 0) > 0.99 ? " · compiling WebGPU shaders" : "") + + "."; + hint.hidden = false; + } +} + +// ------------------------------------------------------------ diagnostics --- + +function renderGpu() { + const el = $("gpu"); + if (navigator.gpu) { + el.hidden = false; + el.textContent = "WebGPU is available in this context."; + return; + } + el.hidden = false; + el.classList.add("error"); + el.textContent = + "navigator.gpu is missing. Set dom.webgpu.enabled = true in about:config (see Firefox setup below) and restart Firefox — models cannot load until then."; +} + +/** Spells out what another engine actually costs, from the loaded model's own record. */ +async function renderPoolCost() { + const { engineCount } = await store.getSettings(); + const record = (await store.list()).find((m) => m.model_id === lastEngineState.modelId) ?? (await store.list())[0]; + if (!record) return void ($("poolCost").textContent = ""); + const weights = record.sizeBytes ?? 0; + const total = weights * engineCount; + // Measured on an M4 Air with a 0.8B model: 2 engines gave 1.6x aggregate + // throughput, 4 gave 0.3x - past the memory budget they starve each other. + const verdict = + engineCount === 1 + ? "no parallelism: batches run one at a time." + : engineCount === 2 + ? "measured ~1.6x aggregate throughput on a 0.8B model." + : "more is usually worse — 4 engines measured 3x SLOWER than 1. Verify with npm run e2e before keeping this."; + $("poolCost").textContent = + `${engineCount} engine(s) x ~${formatBytes(weights)} = ~${formatBytes(total)} VRAM. ${verdict}`; + $("poolCost").classList.toggle("warn", total > 6e9 || engineCount > 2); +} + +/** + * Spells out the sawtooth, because "more steps" is not monotonically better. + * + * Firefox resolves a GPU sync only on a 100 ms tick, so a burst of K steps costs + * a whole number of ticks. The reference figure is the ~7.3 ms/token of real + * compute measured for a 0.8B (AI.md, "The 10 tok/s ceiling"); a bigger model + * costs more per step and wants a smaller K. + */ +const TICK_MS = 100; +const REFERENCE_STEP_MS = 7.3; + +async function renderDecodeCost() { + const { decodeSteps } = await store.getSettings(); + const ticks = Math.ceil((decodeSteps * REFERENCE_STEP_MS) / TICK_MS); + const rate = decodeSteps / ((ticks * TICK_MS) / 1000); + const perTick = Math.floor(TICK_MS / REFERENCE_STEP_MS); + const wastes = decodeSteps > perTick && decodeSteps % perTick !== 0; + $("decodeCost").textContent = + `${decodeSteps} step(s) per sync = ${decodeSteps} token(s) every ${ticks} tick(s) ` + + `≈ ${rate.toFixed(0)} tok/s on a 0.8B (vs 9.6 at 1 step). ` + + (wastes + ? `${decodeSteps} spills past a 100 ms tick boundary — ${perTick} fits inside one tick and measures faster. Confirm with npm run e2e.` + : "Fits the tick grid. Re-check on a larger model: per-step compute grows, so the best K shrinks."); + $("decodeCost").classList.toggle("warn", wastes); +} + +async function renderQuota() { + if (!navigator.storage?.estimate) return; + const { usage, quota } = await navigator.storage.estimate(); + $("quota").textContent = `${formatBytes(usage)} of ${formatBytes(quota)}`; +} + +// --------------------------------------------------------------- ingestion --- + +const drop = $("drop"); + +for (const type of ["dragenter", "dragover"]) { + drop.addEventListener(type, (e) => { + e.preventDefault(); + drop.classList.add("over"); + }); +} +for (const type of ["dragleave", "drop"]) { + drop.addEventListener(type, () => drop.classList.remove("over")); +} +drop.addEventListener("drop", async (e) => { + e.preventDefault(); + await ingest(await filesFromDataTransfer(e.dataTransfer)); +}); + +$("pick").addEventListener("click", () => $("picker").click()); +$("picker").addEventListener("change", async (e) => { + await ingest(filesFromInput(e.target.files)); + e.target.value = ""; +}); + +async function ingest(entries) { + $("ingestError").hidden = true; + $("ingest").hidden = false; + $("ingestBar").style.width = "0%"; + $("ingestStatus").textContent = "Validating…"; + + try { + const record = await ingestModelFolder(entries, { + store, + onProgress: ({ phase, done, total, label }) => { + $("ingestBar").style.width = `${Math.round((done / Math.max(total, 1)) * 100)}%`; + $("ingestStatus").textContent = + phase === "validating" ? label : `Caching ${done}/${total} — ${label}`; + }, + }); + $("ingestStatus").textContent = `Registered ${record.model_id} (${formatBytes(record.sizeBytes)}, ${record.shardCount} shards).`; + await Promise.all([renderModels(), renderQuota()]); + } catch (err) { + $("ingest").hidden = true; + $("ingestError").hidden = false; + $("ingestError").textContent = err.message; + } +} + +// ---------------------------------------------------------------- registry --- + +async function renderModels() { + const models = await store.list(); + const tbody = $("models").querySelector("tbody"); + tbody.replaceChildren(); + $("models").hidden = models.length === 0; + $("noModels").hidden = models.length > 0; + + for (const record of models) { + const tr = document.createElement("tr"); + tr.innerHTML = ` +
+ + + + checking… + `; + + tr.querySelector("strong").textContent = record.model_id; + tr.querySelector("td .muted").textContent = `added ${new Date(record.addedAt).toLocaleString()}`; + const cells = tr.querySelectorAll("td.mono"); + cells[0].textContent = formatBytes(record.sizeBytes); + cells[1].textContent = record.shardCount ?? "—"; + cells[2].textContent = record.wasm ?? "—"; + + tr.querySelector('[data-act="load"]').addEventListener("click", async (e) => { + e.target.disabled = true; + const res = await browser.runtime.sendMessage(request(OP.LOAD, { modelId: record.model_id })); + e.target.disabled = false; + if (!res?.ok) renderEngine({ status: ENGINE_STATE.ERROR, error: res?.error }); + }); + tr.querySelector('[data-act="remove"]').addEventListener("click", async () => { + if (!confirm(`Remove "${record.model_id}" and free ${formatBytes(record.sizeBytes)} of cache?`)) return; + await store.remove(record.model_id); + await Promise.all([renderModels(), renderQuota()]); + }); + + tbody.append(tr); + + store.verify(record).then(({ ok, missing }) => { + const pill = tr.querySelector(".pill"); + pill.classList.add(ok ? "ok" : "bad"); + pill.textContent = ok ? "complete" : `${missing.length} missing`; + }); + } +} + +// ---------------------------------------------------------------- settings --- + +async function renderSettings() { + const s = await store.getSettings(); + $("temperature").value = s.temperature; + $("maxTokens").value = s.maxTokens; + $("engineCount").value = s.engineCount; + $("decodeSteps").value = s.decodeSteps; + $("systemPrompt").value = s.systemPrompt; + $("allowedExternalIds").value = s.allowedExternalIds.join(", "); +} + +$("save").addEventListener("click", async () => { + await store.setSettings({ + temperature: Number($("temperature").value), + maxTokens: Number($("maxTokens").value), + engineCount: Math.max(1, Math.min(4, Number($("engineCount").value) || 1)), + decodeSteps: clampSteps($("decodeSteps").value), + systemPrompt: $("systemPrompt").value, + allowedExternalIds: $("allowedExternalIds").value.split(",").map((s) => s.trim()).filter(Boolean), + }); + $("saved").hidden = false; + setTimeout(() => ($("saved").hidden = true), 1500); + await renderPoolCost(); + await renderDecodeCost(); + renderEngine({ ...lastEngineState, note: "reload the model for the pool size to take effect" }); +}); + +// ------------------------------------------------------------- Engine API --- + +/** + * The three shapes of work, as copy-paste starting points. + * + * They differ only in scheduling metadata, not in op — that is the whole point, + * and the reason there is no `translate` op to call. Each carries the one field + * that makes it behave correctly, because those are what callers get wrong. + */ +function apiSamples(id) { + return { + completion: { + why: "Latency is the product. `session` is what makes this work — reuse one key and the engine drops the stale request itself, so a fast typist never queues a request per keystroke.", + code: `const port = browser.runtime.connect("${id}", { name: "${PORT_NAME}" }); +port.onMessage.addListener((m) => { + if (m.op === "${PORT_OP.CHUNK}") render(m.delta); + if (m.op === "${PORT_OP.DONE}") finish(m.text); +}); + +// On every keystroke. The previous request is superseded, not queued. +port.postMessage({ + protocol: "${PROTOCOL}", + op: "${PORT_OP.CHAT_STREAM}", + id: crypto.randomUUID(), + session: "ghost-text", // supersession key — the field that matters + priority: "${PRIORITY.INTERACTIVE}", // may preempt work that opted in + max_tokens: 24, + messages: [{ role: "user", content: prefix }], +});`, + }, + translation: { + why: "One `batch`, not a loop of `chat` calls: one round trip, and the engine schedules the whole page as a single task that can never occupy more than one engine.", + code: `const res = await browser.runtime.sendMessage("${id}", { + protocol: "${PROTOCOL}", + op: "${OP.BATCH}", + task: "translate-page", // optional; a batch is one task either way + requests: sentences.map((s) => ({ + messages: [{ role: "user", content: \`Translate to French, output only the translation:\\n\${s}\` }], + })), +}); +if (!res.ok) throw new Error(res.error); +res.results.forEach((r) => apply(r.index, r.text));`, + }, + reformat: { + why: "Nobody is watching, so let interactive work cut in. Set `preemptible` on the job that can afford to lose — a preempted job resolves with `preempted: true` and its partial text, never requeued.", + code: `const res = await browser.runtime.sendMessage("${id}", { + protocol: "${PROTOCOL}", + op: "${OP.CHAT}", + priority: "${PRIORITY.BACKGROUND}", + preemptible: true, // the direction matters + max_tokens: 2048, + messages: [{ role: "user", content: \`Reformat as clean Markdown, no commentary:\\n\\n\${doc}\` }], +}); +if (res.preempted) keepOrDiscard(res.text);`, + }, + }; +} + +function renderApiSample() { + const id = browser.runtime.id; + const samples = apiSamples(id); + $("selfId").textContent = id; + + const show = (name) => { + for (const tab of document.querySelectorAll(".tab")) { + tab.classList.toggle("active", tab.dataset.tab === name); + } + $("apiWhy").textContent = samples[name].why; + $("apiSample").textContent = samples[name].code; + }; + + for (const tab of document.querySelectorAll(".tab")) { + tab.addEventListener("click", () => show(tab.dataset.tab)); + } + show("completion"); + + const flash = (btn, label) => { + const original = btn.textContent; + btn.textContent = label; + setTimeout(() => (btn.textContent = original), 1200); + }; + $("copyId").addEventListener("click", () => + navigator.clipboard.writeText(id).then(() => flash($("copyId"), "Copied")), + ); + $("copyCode").addEventListener("click", () => + navigator.clipboard.writeText($("apiSample").textContent).then(() => flash($("copyCode"), "Copied")), + ); +} + +renderGpu(); +renderApiSample(); +await Promise.all([renderModels(), renderSettings(), renderQuota()]); +await renderPoolCost(); +await renderDecodeCost(); +browser.runtime.sendMessage(request(OP.STATUS)).then((res) => res?.ok && renderEngine(res.state)); diff --git a/src/popup/popup.css b/src/popup/popup.css new file mode 100644 index 0000000000000000000000000000000000000000..21d9f600dbc1eddadf52f08952938aa2241ec89a --- /dev/null +++ b/src/popup/popup.css @@ -0,0 +1,50 @@ +body { width: 420px; height: 580px; display: flex; flex-direction: column; } + +header, footer { padding: 12px 14px; display: flex; flex-direction: column; gap: 10px; } +header { border-bottom: 1px solid var(--line); } +footer { border-top: 1px solid var(--line); gap: 8px; } + +main { + flex: 1; + overflow-y: auto; + padding: 16px 14px; + display: flex; + flex-direction: column; + gap: 18px; +} + +/* Claude does not put the assistant in a bubble: the reply is just text on the + page, and only the user's turn gets a surface. Keeping that asymmetry is most + of what makes a transcript read as Claude's rather than as a chat app's. */ +.msg { white-space: pre-wrap; word-break: break-word; } +.msg.user { + align-self: flex-end; + max-width: 85%; + background: var(--panel-strong); + border-radius: 12px; + padding: 9px 13px; +} +.msg.assistant { align-self: stretch; } +.msg.error { + align-self: stretch; + border: 1px solid var(--accent); + border-radius: 8px; + padding: 9px 13px; +} +.msg .meta { display: block; margin-top: 6px; font-size: 11px; color: var(--muted); } + +.empty { margin: auto; text-align: center; color: var(--muted); max-width: 260px; } + +#input { + resize: none; + width: 100%; + background: var(--panel); + border-radius: 12px; + padding: 10px 12px; +} +#input:focus { border-color: var(--line-strong); } + +#status { min-height: 16px; overflow: hidden; text-overflow: ellipsis; white-space: nowrap; font-size: 12px; } +/* The load report is a sentence, not a label: let it wrap while it is showing, + otherwise it is silently truncated to "Loading model from cache[26/5…". */ +#status.loading { white-space: normal; overflow: visible; text-overflow: clip; } diff --git a/src/popup/popup.html b/src/popup/popup.html new file mode 100644 index 0000000000000000000000000000000000000000..a3bbe0f0bab0b5617b950e7a5437babb5a912fee --- /dev/null +++ b/src/popup/popup.html @@ -0,0 +1,40 @@ + + + + + Everything WebGPU + + + + +
+
+ +

Everything WebGPU

+ +
+
+ + + +
+
+
Checking engine…
+ +
+ +
+ +
+ +
+ + + + +
+
+ + + + diff --git a/src/popup/popup.js b/src/popup/popup.js new file mode 100644 index 0000000000000000000000000000000000000000..093183c0585f8737906a9a99aa3417a99765c85c --- /dev/null +++ b/src/popup/popup.js @@ -0,0 +1,235 @@ +/** Minimal test chat. Thin client over the background engine host's port protocol. */ +import { ENGINE_STATE, OP, PORT_NAME, PORT_OP, PRIORITY, PROTOCOL, request } from "../adapters/protocol.js"; + +const $ = (id) => document.getElementById(id); +const els = { + dot: $("dot"), model: $("model"), load: $("load"), unload: $("unload"), + progress: $("progress"), status: $("status"), loadHint: $("loadHint"), log: $("log"), input: $("input"), + send: $("send"), stop: $("stop"), clear: $("clear"), manage: $("manage"), +}; + +const history = []; +let models = []; +let engineState = { status: ENGINE_STATE.IDLE, modelId: null, pool: null }; +let streamId = null; +let streamEl = null; + +/** + * The engine supersedes by `session`, so a cleared conversation has to become a + * *different* session — otherwise the next message looks to the scheduler like + * a continuation of the one just thrown away. + */ +const newSession = () => `popup-${crypto.randomUUID()}`; +let session = newSession(); + +const port = browser.runtime.connect({ name: PORT_NAME }); +port.onMessage.addListener(onPortMessage); + +// -------------------------------------------------------------- rendering --- + +/** + * The second line of the load guide. + * + * Deliberately not WebLLM's stock wording. Its default hint is about populating + * a cache from the network on first visit, which is never what happens here: + * the weights were injected into Cache Storage by drag-and-drop, so nothing is + * downloaded and a "later refreshes are faster" promise would be wrong. What + * the user is actually waiting on, once the shards are read, is WebGPU shader + * compilation. + */ +function renderLoadHint(progress) { + if (!progress) return void (els.loadHint.hidden = true); + const secs = progress.timeElapsed; + const parts = []; + if (progress.engines > 1) parts.push(`${progress.engines} engines, loaded 2 at a time`); + parts.push("weights come from the local cache — nothing is downloaded"); + if ((progress.progress ?? 0) > 0.99) { + parts.push("compiling WebGPU shaders, which is most of the wait"); + } + els.loadHint.textContent = `${secs ? `${secs}s elapsed · ` : ""}${parts.join(" · ")}.`; + els.loadHint.hidden = false; +} + +function renderState() { + const { status, modelId, progress, error, pool } = engineState; + els.dot.className = `dot ${status}`; + els.progress.style.width = `${Math.round((progress?.progress ?? (status === ENGINE_STATE.READY ? 1 : 0)) * 100)}%`; + + els.status.textContent = + error ? error : + progress?.text ? progress.text : + status === ENGINE_STATE.READY ? `Ready — ${modelId}` : + status === ENGINE_STATE.LOADING ? `Loading ${modelId}…` : + "Idle — pick a model and press Load"; + if (!error && status === ENGINE_STATE.READY && pool && (pool.busy || pool.queued)) { + const cap = pool.size < pool.maxSize && !pool.growthBlocked ? `/${pool.maxSize}` : ""; + els.status.textContent += ` · ${pool.busy}/${pool.size}${cap} busy, ${pool.queued} queued`; + } + els.status.classList.toggle("banner", Boolean(error)); + els.status.classList.toggle("error", Boolean(error)); + + const loading = status === ENGINE_STATE.LOADING; + // The report is a full sentence from WebLLM ("Loading model from cache[26/58]: + // 890MB loaded. 51% completed, 121 secs elapsed."). #status is a single + // ellipsised line by default, which cut it off, so unclip it while loading. + els.status.classList.toggle("loading", loading && Boolean(progress?.text)); + renderLoadHint(loading ? progress : null); + els.load.disabled = loading || !els.model.value; + els.unload.disabled = loading || status !== ENGINE_STATE.READY; + els.model.disabled = loading; + els.send.disabled = loading || streamId !== null; + els.stop.hidden = streamId === null; + els.load.textContent = status === ENGINE_STATE.READY && modelId === els.model.value ? "Reload" : "Load"; +} + +function renderEmpty() { + els.log.replaceChildren( + Object.assign(document.createElement("div"), { + className: "empty", + textContent: models.length + ? "Load a model, then send a message." + : "No local models yet. Open Models… and drop a compiled MLC model folder.", + }), + ); +} + +function addMessage(role, text) { + if (els.log.querySelector(".empty")) els.log.replaceChildren(); + const el = document.createElement("div"); + el.className = `msg ${role}`; + el.textContent = text; + els.log.append(el); + els.log.scrollTop = els.log.scrollHeight; + return el; +} + +// ---------------------------------------------------------------- actions --- + +async function ask(op, payload) { + const res = await browser.runtime.sendMessage(request(op, payload)); + if (!res?.ok) throw new Error(res?.error ?? "Engine did not respond."); + return res; +} + +async function refreshModels() { + const res = await ask(OP.LIST_MODELS); + models = res.models; + engineState = res.state; + els.model.replaceChildren( + ...models.map((m) => new Option(`${m.model_id} · ${(m.sizeBytes / 1e9).toFixed(1)} GB`, m.model_id)), + ); + if (res.state.modelId) els.model.value = res.state.modelId; + if (!els.log.children.length || els.log.querySelector(".empty")) renderEmpty(); + renderState(); +} + +function send() { + const text = els.input.value.trim(); + if (!text) return; + els.input.value = ""; + history.push({ role: "user", content: text }); + addMessage("user", text); + + streamId = crypto.randomUUID(); + streamEl = addMessage("assistant", ""); + renderState(); + + port.postMessage({ + protocol: PROTOCOL, + op: PORT_OP.CHAT_STREAM, + id: streamId, + modelId: els.model.value || undefined, + messages: history, + priority: PRIORITY.INTERACTIVE, + session, + }); +} + +function onPortMessage(msg) { + if (msg?.protocol !== PROTOCOL) return; + switch (msg.op) { + case PORT_OP.ENGINE_STATE: + engineState = msg.state; + if (msg.state.modelId && !els.model.value) els.model.value = msg.state.modelId; + renderState(); + return; + case PORT_OP.CHUNK: + if (msg.id !== streamId) return; + streamEl.textContent += msg.delta; + els.log.scrollTop = els.log.scrollHeight; + return; + case PORT_OP.DONE: { + if (msg.id !== streamId) return; + const text = streamEl.textContent; + history.push({ role: "assistant", content: text }); + if (msg.usage) { + const meta = document.createElement("span"); + meta.className = "meta"; + meta.textContent = `${msg.usage.completion_tokens} tok · ${(msg.usage.extra?.decode_tokens_per_s ?? 0).toFixed(1)} tok/s`; + streamEl.append(meta); + } + streamId = null; + streamEl = null; + renderState(); + return; + } + case PORT_OP.ERROR: + if (streamEl && !streamEl.textContent) streamEl.remove(); + addMessage("error", msg.error); + streamId = null; + streamEl = null; + renderState(); + return; + } +} + +// ------------------------------------------------------------------ wiring --- + +els.load.addEventListener("click", () => + ask(OP.LOAD, { modelId: els.model.value }).catch((err) => { + engineState = { ...engineState, status: ENGINE_STATE.ERROR, error: err.message }; + renderState(); + }), +); +els.unload.addEventListener("click", () => ask(OP.UNLOAD).catch(() => {})); +els.send.addEventListener("click", send); +els.stop.addEventListener("click", () => + port.postMessage({ protocol: PROTOCOL, op: PORT_OP.ABORT, id: streamId, session }), +); + +/** + * Discard the conversation and start a fresh session. + * + * All three parts matter. Dropping `history` alone left every rendered message + * on screen, which read as the button doing nothing; and leaving the old + * session key in place meant a request that was still streaming kept writing + * into a conversation the user had already thrown away. + */ +els.clear.addEventListener("click", () => { + if (streamId) { + port.postMessage({ protocol: PROTOCOL, op: PORT_OP.ABORT, id: streamId, session }); + streamId = null; + streamEl = null; + } + history.length = 0; + session = newSession(); + renderEmpty(); + els.input.value = ""; + renderState(); +}); +els.manage.addEventListener("click", () => { + browser.runtime.openOptionsPage(); + window.close(); +}); +els.model.addEventListener("change", renderState); +els.input.addEventListener("keydown", (e) => { + if (e.key === "Enter" && !e.shiftKey) { + e.preventDefault(); + if (!els.send.disabled) send(); + } +}); + +refreshModels().catch((err) => { + els.status.textContent = err.message; + els.status.classList.add("banner", "error"); +}); diff --git a/src/ui.css b/src/ui.css new file mode 100644 index 0000000000000000000000000000000000000000..8946aef0ba7d706ea621707eb7d18a906f71182d --- /dev/null +++ b/src/ui.css @@ -0,0 +1,108 @@ +/** + * Shared UI tokens, in Claude's palette. + * + * Exactly three colours are declared per theme — a warm off-white ground, a + * near-black ink, and Claude's coral accent. Everything that looks like a + * fourth colour (borders, muted text, panels) is one of those three at reduced + * alpha, so the whole surface stays in one family and nothing needs picking a + * new hue to add a component. + * + * Fonts are a system stack. Claude's own Styrene and Tiempos are licensed + * faces, and this extension does no network requests by design, so there is + * nothing to load them from — the stack below is the closest available match. + */ +:root { + --bg: #faf9f5; + --ink: #141413; + --accent: #d97757; + + --panel: color-mix(in srgb, var(--ink) 4%, var(--bg)); + --panel-strong: color-mix(in srgb, var(--ink) 7%, var(--bg)); + --line: color-mix(in srgb, var(--ink) 10%, transparent); + --line-strong: color-mix(in srgb, var(--ink) 18%, transparent); + --muted: color-mix(in srgb, var(--ink) 52%, transparent); + --on-accent: #ffffff; + + color-scheme: light; +} + +@media (prefers-color-scheme: dark) { + :root { + --bg: #262624; + --ink: #f5f4ef; + --accent: #d97757; + + --panel: color-mix(in srgb, var(--ink) 6%, var(--bg)); + --panel-strong: color-mix(in srgb, var(--ink) 10%, var(--bg)); + --line: color-mix(in srgb, var(--ink) 14%, transparent); + --line-strong: color-mix(in srgb, var(--ink) 24%, transparent); + --muted: color-mix(in srgb, var(--ink) 55%, transparent); + + color-scheme: dark; + } +} + +* { box-sizing: border-box; } + +body { + margin: 0; + background: var(--bg); + color: var(--ink); + font: 14px/1.6 ui-sans-serif, -apple-system, BlinkMacSystemFont, "SF Pro Text", "Segoe UI", system-ui, sans-serif; + -webkit-font-smoothing: antialiased; +} + +h1, h2 { margin: 0; font-weight: 500; letter-spacing: -0.01em; } +h1 { font-size: 15px; } +h2 { font-size: 12px; color: var(--muted); font-weight: 500; letter-spacing: .04em; text-transform: uppercase; } + +button, select, input, textarea { + font: inherit; + color: var(--ink); + background: transparent; + border: 1px solid var(--line-strong); + border-radius: 8px; + padding: 7px 12px; + transition: background .12s ease, border-color .12s ease, opacity .12s ease; +} +select { background: var(--panel); } +button { cursor: pointer; } +button:hover:not(:disabled) { background: var(--panel-strong); } +button:disabled { opacity: .35; cursor: default; } + +button.primary { + background: var(--accent); + color: var(--on-accent); + border-color: var(--accent); + font-weight: 500; +} +button.primary:hover:not(:disabled) { background: color-mix(in srgb, var(--ink) 12%, var(--accent)); } +/* Destructive actions read as plain text until hovered — Claude does not spend + a colour on them, and there is no fourth colour to spend. */ +button.danger { border-color: transparent; color: var(--muted); } +button.danger:hover:not(:disabled) { color: var(--accent); background: var(--panel-strong); } + +:focus-visible { outline: 2px solid var(--accent); outline-offset: 1px; } + +.row { display: flex; gap: 8px; align-items: center; } +.grow { flex: 1; min-width: 0; } +.muted { color: var(--muted); } +.mono { font-family: ui-monospace, "SF Mono", Menlo, monospace; font-size: 12px; } + +.dot { width: 7px; height: 7px; border-radius: 50%; background: var(--line-strong); flex: none; } +.dot.ready { background: var(--accent); } +.dot.loading { background: var(--accent); animation: pulse 1.4s ease-in-out infinite; } +.dot.error { background: var(--ink); } +@keyframes pulse { 50% { opacity: .3; } } + +.hint { font-size: 12px; line-height: 1.45; margin-top: 3px; } +.bar { height: 2px; background: var(--line); border-radius: 2px; overflow: hidden; } +.bar > i { display: block; height: 100%; background: var(--accent); width: 0; transition: width .2s; } + +.banner { + padding: 10px 12px; + border-radius: 8px; + border: 1px solid var(--line-strong); + background: var(--panel); +} +.banner.error { border-color: var(--accent); } diff --git a/src/vite.js b/src/vite.js new file mode 100644 index 0000000000000000000000000000000000000000..21e13d34ee15954c3e592a622d53b84cf98d96cb --- /dev/null +++ b/src/vite.js @@ -0,0 +1,71 @@ +/** + * The Vite plugin, because one line of config you have to *remember* is one + * line of config you will forget. + * + * ```js + * import { everythingWebGPU } from "everything-webgpu/vite"; + * export default defineConfig({ plugins: [everythingWebGPU()] }); + * ``` + * + * ## What it fixes, precisely + * + * The engine spawns its decode worker with + * `new Worker(new URL("./engine-worker.js", import.meta.url), { type: "module" })`. + * Vite's dependency pre-bundler runs esbuild over `node_modules` packages and + * copies that expression through **verbatim** — but the copy now lives in + * `node_modules/.vite/deps/everything-webgpu.js`, so `import.meta.url` points + * there and the worker resolves to `.vite/deps/engine-worker.js`, which does not + * exist. The real file is still at + * `node_modules/everything-webgpu/src/engine/engine-worker.js`. + * + * Measured, not assumed. On a real (non-linked) install of this package: + * + * | | worker resolves | notes | + * | --- | --- | --- | + * | `vite build` | ✅ | Rollup handles it; output is byte-identical with or without this plugin | + * | `vite dev`, excluded | ✅ | source served from `/@fs/`, `import.meta.url` is correct | + * | `vite dev`, not excluded | ❌ 404 | the case this exists for | + * + * Two consequences worth stating, because both mislead: + * + * 1. **`vite build` never reproduces the bug.** A green production build is + * not evidence that the dev server works. + * 2. **A linked (`file:`) dependency never reproduces it either**, because + * Vite does not pre-bundle linked packages. Every example in this repo is + * linked, so none of them can catch this — only a real install can. See + * `npm run verify-consumer`. + * + * The lazy `import("../../vendor/web-llm.js")` is *not* affected: esbuild + * rewrites that one correctly to a hashed chunk. Only the worker breaks. + * + * If you would rather not add a plugin, the equivalent is: + * + * ```js + * optimizeDeps: { exclude: ["everything-webgpu"] } + * ``` + * + * and if you do neither, `load()` fails with `PACKAGE_INCOMPLETE` naming this + * fix rather than hanging on a worker that will never answer. + */ + +/** The package's own name, so the exclusion cannot drift from it. */ +const PACKAGE = "everything-webgpu"; + +/** + * @param {{exclude?: string[]}} [opts] `exclude` adds further specifiers, for a + * host that re-exports this engine from its own package and hits the same + * pre-bundling of the worker URL. + * @returns {import("vite").Plugin} + */ +export function everythingWebGPU({ exclude = [] } = {}) { + return { + name: "everything-webgpu", + // `config` rather than `configResolved`: this has to merge into the user's + // options before Vite computes the optimizer's entries, not after. + config: () => ({ + optimizeDeps: { exclude: [PACKAGE, ...exclude] }, + }), + }; +} + +export default everythingWebGPU; diff --git a/tsconfig.json b/tsconfig.json new file mode 100644 index 0000000000000000000000000000000000000000..d7f0df07bddd3a6e5ffdcda2e35834b58f6f3166 --- /dev/null +++ b/tsconfig.json @@ -0,0 +1,31 @@ +{ + "//": "Type declarations only. The source stays plain ESM JavaScript; the JSDoc it already carries is the type source, so `npm run types` adds a typed surface for consumers without a rewrite or a build step on the runtime path.", + "compilerOptions": { + "allowJs": true, + "checkJs": false, + "declaration": true, + "emitDeclarationOnly": true, + "outDir": "types", + "target": "ES2022", + "module": "ESNext", + "moduleResolution": "bundler", + "lib": [ + "ES2022", + "DOM", + "DOM.Iterable", + "WebWorker" + ], + "skipLibCheck": true, + "strict": false, + "rootDir": "src" + }, + "include": [ + "src/engine/**/*.js", + "src/adapters/**/*.js", + "src/vite.js" + ], + "exclude": [ + "src/engine/engine-worker.js", + "vendor" + ] +} diff --git a/types/adapters/idb.d.ts b/types/adapters/idb.d.ts new file mode 100644 index 0000000000000000000000000000000000000000..a9bac14d9cc2761ce4c92dc9f681f3a7c57dfe02 --- /dev/null +++ b/types/adapters/idb.d.ts @@ -0,0 +1,18 @@ +/** @returns {Promise} */ +export function indexedDBStorage(): Promise; +/** + * Ask the browser to stop treating this origin's storage as evictable. + * + * Firefox prompts the user (or grants silently for a site with the permission + * already); Chrome decides from engagement heuristics without a prompt. Either + * way a caller must handle `persisted: false` — the model still works, it can + * just be dropped under disk pressure, which `ModelStore.verify()` will catch + * on the next load. + * + * @returns {Promise<{persisted: boolean, quota?: number, usage?: number}>} + */ +export function ensurePersistent(): Promise<{ + persisted: boolean; + quota?: number; + usage?: number; +}>; diff --git a/types/adapters/memory.d.ts b/types/adapters/memory.d.ts new file mode 100644 index 0000000000000000000000000000000000000000..25f84af8cba7134201f0ae450ba77f4f0dd780ea --- /dev/null +++ b/types/adapters/memory.d.ts @@ -0,0 +1,10 @@ +/** + * In-memory StorageAdapter. + * + * For tests, and for a host that genuinely wants the registry to die with the + * page. Note that the *weights* still live in Cache Storage and survive — only + * the registry entry pointing at them is lost, which `ModelStore.verify()` will + * then never be asked about. Do not use this in production for that reason. + */ +/** @returns {import("../engine/model-store.js").StorageAdapter} */ +export function memoryStorage(initial?: {}): import("../engine/model-store.js").StorageAdapter; diff --git a/types/adapters/protocol.d.ts b/types/adapters/protocol.d.ts new file mode 100644 index 0000000000000000000000000000000000000000..34dd54a071baa0864f94c39204b12d8b7335f600 --- /dev/null +++ b/types/adapters/protocol.d.ts @@ -0,0 +1,28 @@ +export function request(op: any, payload?: {}): { + protocol: string; + op: any; +}; +export const PROTOCOL: "everything-webgpu/v1"; +export const PORT_NAME: "everything-webgpu/v1"; +export namespace OP { + let STATUS: string; + let LIST_MODELS: string; + let LOAD: string; + let UNLOAD: string; + let CHAT: string; + let BATCH: string; + let CANCEL: string; + let CONFIGURE: string; +} +export namespace PORT_OP { + let CHAT_STREAM: string; + let BATCH_STREAM: string; + let ABORT: string; + let SUBSCRIBE: string; + let CHUNK: string; + let ITEM: string; + let DONE: string; + let ERROR: string; + let ENGINE_STATE: string; +} +export { ENGINE_STATE, PRIORITY, PRIORITY_ORDER } from "../engine/constants.js"; diff --git a/types/adapters/webext.d.ts b/types/adapters/webext.d.ts new file mode 100644 index 0000000000000000000000000000000000000000..bf74009a854e82c07a9f3e2e70bdd75100d782d0 --- /dev/null +++ b/types/adapters/webext.d.ts @@ -0,0 +1,15 @@ +/** + * `browser.storage.local` already *is* the StorageAdapter shape — `get(key)` + * and `set(obj)`. The two-method interface was chosen for that reason, so this + * is a passthrough rather than a translation layer. + * + * @returns {import("../engine/model-store.js").StorageAdapter} + */ +export function webExtensionStorage(): import("../engine/model-store.js").StorageAdapter; +/** + * Wires a ScheduledEngine onto this extension's runtime messaging. + * + * @param {import("../engine/engine.js").ScheduledEngine} engine + * @returns {() => void} detach + */ +export function attachWebExtensionTransport(engine: import("../engine/engine.js").ScheduledEngine): () => void; diff --git a/types/engine/chat.d.ts b/types/engine/chat.d.ts new file mode 100644 index 0000000000000000000000000000000000000000..e8fcd0cd09034d8bf9548187195f81559827a826 --- /dev/null +++ b/types/engine/chat.d.ts @@ -0,0 +1,18 @@ +/** @param {import("./engine.js").ScheduledEngine} engine */ +export function chatFacade(engine: import("./engine.js").ScheduledEngine): { + completions: { + /** + * @param {import("./engine.js").CompletionRequest & { + * stream?: boolean, stream_options?: {include_usage?: boolean} }} req + * @returns {Promise>} a completion, or a + * stream of chunks when `stream` is set — the same two shapes WebLLM + * returns, so `await`ing then `for await`ing works unchanged. + */ + create(req: import("./engine.js").CompletionRequest & { + stream?: boolean; + stream_options?: { + include_usage?: boolean; + }; + }): Promise>; + }; +}; diff --git a/types/engine/constants.d.ts b/types/engine/constants.d.ts new file mode 100644 index 0000000000000000000000000000000000000000..d129fde0c690b5ffee69dad2e8b18ab9346ab66a --- /dev/null +++ b/types/engine/constants.d.ts @@ -0,0 +1,26 @@ +export namespace PRIORITY { + let INTERACTIVE: string; + let NORMAL: string; + let BACKGROUND: string; +} +export const PRIORITY_ORDER: string[]; +export namespace UNLOAD_LEVEL { + let VRAM: string; + let CACHE: string; +} +export namespace JOB_KIND { + let CHAT: string; + let EMBEDDING: string; +} +export namespace ENGINE_STATE { + let IDLE: string; + let LOADING: string; + let READY: string; + let ERROR: string; +} +/** + * Out-of-band message kind for configuring an engine worker before WebLLM's own + * handshake starts. Distinct from every WebLLM `kind`, so the worker can route + * on it without parsing the rest. + */ +export const WORKER_CONFIGURE: "everything-webgpu/configure"; diff --git a/types/engine/create.d.ts b/types/engine/create.d.ts new file mode 100644 index 0000000000000000000000000000000000000000..73e73443f13dd7e0eed7b69f7fd5c2890aa797f3 --- /dev/null +++ b/types/engine/create.d.ts @@ -0,0 +1,24 @@ +/** + * @param {string} [modelId] loaded before returning, as `CreateMLCEngine` does. + * Omit to get an engine that loads later. + * @param {object} [opts] + * @param {ModelStore | import("./model-store.js").StorageAdapter} [opts.store] + * defaults to IndexedDB. Pass one explicitly in a worker or a test. + * @param {(report: {text: string, progress: number}) => void} [opts.initProgressCallback] + * @param {string | URL} [opts.workerUrl] + * @param {() => Promise} [opts.loadWebLLM] + * @param {boolean} [opts.prebuilt] + * @returns {Promise} + */ +export function CreateScheduledEngine(modelId?: string, opts?: { + store?: ModelStore | import("./model-store.js").StorageAdapter; + initProgressCallback?: (report: { + text: string; + progress: number; + }) => void; + workerUrl?: string | URL; + loadWebLLM?: () => Promise; + prebuilt?: boolean; +}): Promise; +import { ModelStore } from "./model-store.js"; +import { ScheduledEngine } from "./engine.js"; diff --git a/types/engine/device.d.ts b/types/engine/device.d.ts new file mode 100644 index 0000000000000000000000000000000000000000..819445c71a90bbd95e17765187c03a5418bbf476 --- /dev/null +++ b/types/engine/device.d.ts @@ -0,0 +1,142 @@ +/** + * @typedef {object} DeviceProbe + * @property {boolean} webgpu + * @property {string} [reason] why WebGPU is unusable, when it is + * @property {object} [adapter] vendor / architecture / device, where exposed + * @property {{shaderF16: boolean}} [features] + * @property {object} [limits] + * @property {boolean} [kvReuse] whether cross-turn KV reuse can be used here + * @property {{quota?: number, usage?: number, persisted?: boolean}} storage + * @property {number} [deviceMemoryGB] Chrome only; absent is not "small" + */ +/** + * Reads what this machine will admit to. Never throws: an unusable device is a + * result, not an error — the caller's job is to explain it, not to crash. + * + * @returns {Promise} + */ +export function probeDevice(): Promise; +/** + * Whether a model can run on a probed device. + * + * `blockers` mean it will not work; `warnings` mean it will work worse, or + * might not fit. The split matters: a caller should refuse to start on a + * blocker and merely inform on a warning, and conflating the two is how you end + * up either crashing or refusing to run something that would have been fine. + * + * @param {{model_id: string, vram_required_MB?: number, sizeBytes?: number}} model + * @param {DeviceProbe} probe + * @returns {{ok: boolean, blockers: Array<{code: string, message: string}>, + * warnings: Array<{code: string, message: string}>}} + */ +export function canRun(model: { + model_id: string; + vram_required_MB?: number; + sizeBytes?: number; +}, probe: DeviceProbe): { + ok: boolean; + blockers: Array<{ + code: string; + message: string; + }>; + warnings: Array<{ + code: string; + message: string; + }>; +}; +/** + * Rank a model list by what this device can actually run. + * + * The prebuilt list spans 239 MB to 31 GB, so "which model should I use" is the + * first question a developer has and the one they have least basis to answer. + * Runnable models come first, then fewest warnings; unrunnable ones are kept at + * the end carrying their reason rather than silently dropped, because "why + * can't I use that one" is the next question. + * + * **`prefer` is a real choice, not a default worth hiding.** Decode here is + * memory-bandwidth-bound — time per token scales with weight bytes (AI.md, + * "Why not llama.cpp/Ollama-class"), so the largest model that fits is also the + * slowest thing that fits. `"quality"` picks the biggest, `"speed"` the + * smallest. Neither is right for everyone, which is why it is a parameter. + * + * Vision models are excluded from a text ranking rather than merely deprioritised: + * a VLM answers text prompts perfectly well, but at several times the download + * for no benefit, so recommending one to a caller who did not ask is bad advice. + * + * @param {Array} models `model_list` entries or registry records + * @param {{probe: DeviceProbe, maxVramMB?: number, needsVision?: boolean, + * prefer?: "quality" | "speed"}} opts + */ +export function rankModels(models: Array, { probe, maxVramMB, needsVision, prefer }?: { + probe: DeviceProbe; + maxVramMB?: number; + needsVision?: boolean; + prefer?: "quality" | "speed"; +}): { + ok: boolean; + blockers: Array<{ + code: string; + message: string; + }>; + warnings: Array<{ + code: string; + message: string; + }>; + model: any; +}[]; +/** + * @param {number} modelBytes + * @param {number} [bytesPerSecond] this machine's measured rate, when known + * @returns {{tokensPerSecond: number, basis: "measured" | "extrapolated", + * modelBytes: number, bytesPerSecond: number, reference?: string}} + */ +export function projectSpeed(modelBytes: number, bytesPerSecond?: number): { + tokensPerSecond: number; + basis: "measured" | "extrapolated"; + modelBytes: number; + bytesPerSecond: number; + reference?: string; +}; +/** + * Decode throughput is memory bandwidth divided by weight bytes. + * + * This project measured the whole chain: decode reaches ~16 GB/s of the M4's + * ~120 GB/s, and the 1.06 GB build runs 16.6–18.1 tok/s — which is that + * quotient. So a projection needs one number, the *achieved* bandwidth, and + * everything else follows from model size. + * + * The constant below is that machine's figure and is only a starting point. The + * moment this engine has decoded anything it knows the real number for the + * machine it is on, and `ScheduledEngine.estimateSpeed()` switches to it — so + * this is a cold-start default, not a claim about anyone's hardware. + */ +export const REFERENCE_DECODE_BYTES_PER_SECOND: 17000000000; +export const REFERENCE_DEVICE: "M4 MacBook Air (16 GB), Firefox"; +export type DeviceProbe = { + webgpu: boolean; + /** + * why WebGPU is unusable, when it is + */ + reason?: string; + /** + * vendor / architecture / device, where exposed + */ + adapter?: object; + features?: { + shaderF16: boolean; + }; + limits?: object; + /** + * whether cross-turn KV reuse can be used here + */ + kvReuse?: boolean; + storage: { + quota?: number; + usage?: number; + persisted?: boolean; + }; + /** + * Chrome only; absent is not "small" + */ + deviceMemoryGB?: number; +}; diff --git a/types/engine/engine.d.ts b/types/engine/engine.d.ts new file mode 100644 index 0000000000000000000000000000000000000000..3d08a26a7bedc4f33758a46222f20e9454d995ea --- /dev/null +++ b/types/engine/engine.d.ts @@ -0,0 +1,582 @@ +export class ScheduledEngine { + /** + * @param {object} opts + * @param {ModelStore | import("./model-store.js").StorageAdapter} opts.store + * a ModelStore, or a bare StorageAdapter to wrap in one + * @param {string | URL} [opts.workerUrl] + * @param {() => Promise} [opts.loadWebLLM] + * @param {boolean} [opts.prebuilt] expose WebLLM's 163 HuggingFace-hosted + * models, downloaded on first load. Default true. Set false for an + * offline-only build: `load()` then resolves registered models and nothing + * else, and an unknown id fails before the WebLLM bundle is even fetched. + */ + constructor({ store, workerUrl, loadWebLLM, prebuilt }?: { + store: ModelStore | import("./model-store.js").StorageAdapter; + workerUrl?: string | URL; + loadWebLLM?: () => Promise; + prebuilt?: boolean; + }); + /** The ModelStore, so a host can drive the registry without a second handle. */ + get store(): ModelStore; + /** + * `chat.completions.create()`, the WebLLM/OpenAI shape. See `chat.js`. + * + * Built once and cached: callers hold on to `engine.chat.completions` the way + * they did with WebLLM, and a fresh object each access would break that. + */ + get chat(): any; + /** + * `environment()` — the read-only report, with `environment.measure()` on it. + * + * Cached like `chat` so a caller can hold on to it. Writes are `configure()`; + * see `environment.js` for why those are separate verbs. + */ + get environment(): any; + get state(): { + status: string; + modelId: any; + progress: any; + error: any; + pool: { + size: number; + busy: number; + queued: number; + }; + /** Model ids with a live pool. `modelId` is whichever of them is current. */ + resident: any[]; + /** Latest decode probe from an engine worker; see multistep.js. */ + decode: any; + }; + get hasWebGPU(): boolean; + /** + * @param {(state: object) => void} listener called immediately, then on change + * @returns {() => void} unsubscribe + */ + subscribe(listener: (state: object) => void): () => void; + /** Model ids with a live pool right now. */ + get resident(): string[]; + /** + * Choose which resident model unaddressed requests go to. + * + * Distinct from `load()` on purpose: this is free and instant, because the + * weights are already up. `load()` is what costs. + */ + use(modelId: any): { + status: string; + modelId: any; + progress: any; + error: any; + pool: { + size: number; + busy: number; + queued: number; + }; + /** Model ids with a live pool. `modelId` is whichever of them is current. */ + resident: any[]; + /** Latest decode probe from an engine worker; see multistep.js. */ + decode: any; + }; + /** Registered models only — cheap, no bundle load. */ + listModels(): Promise; + /** + * Everything `load()` would accept, normalised: registered models first, then + * WebLLM's prebuilt list. + * + * Costs a WebLLM bundle fetch when `prebuilt` is on, because the list lives + * inside it. `listModels()` is the cheap call if you only care about what this + * app registered. + * + * @returns {Promise>} + */ + listAvailableModels(): Promise>; + /** + * What this machine will admit to: WebGPU, adapter, `shader-f16`, the five + * limits that matter, storage quota. Cached — hardware does not change + * mid-session, and `requestAdapter()` is not free. + * @returns {Promise} + */ + probe(): Promise; + /** + * Whether a model will run here, before anything is downloaded. + * @param {string} modelId + * @returns {Promise<{ok: boolean, blockers: Array, warnings: Array}>} + */ + canRun(modelId: string): Promise<{ + ok: boolean; + blockers: Array; + warnings: Array; + }>; + /** + * Which models this device should actually be asked to run, best first. + * + * The prebuilt list spans 239 MB to 31 GB; this is the answer to the first + * question a developer has and the one they have least basis to answer. + * + * @param {{maxVramMB?: number, needsVision?: boolean, needsToolCalling?: boolean, + * prefer?: "quality" | "speed"}} [opts] + */ + recommendModels({ needsToolCalling, ...opts }?: { + maxVramMB?: number; + needsVision?: boolean; + needsToolCalling?: boolean; + prefer?: "quality" | "speed"; + }): Promise<{ + ok: boolean; + blockers: Array<{ + code: string; + message: string; + }>; + warnings: Array<{ + code: string; + message: string; + }>; + model: any; + }[]>; + /** + * Is this model's data on disk, so a load would need no network? + * + * Routes by who knows the keys. We wrote an injected model's artifacts and + * hold the manifest, so `verify()` answers exactly — including a `"partial"` + * verdict WebLLM cannot give. Everything else was fetched by WebLLM, which + * derives the keys as its loader did, so `hasModelInCache` is the answer. + * + * @returns {Promise<"cached" | "partial" | "absent">} + */ + cacheState(modelId: any): Promise<"cached" | "partial" | "absent">; + /** + * Download a model into the cache **without building an engine**. + * + * For warming during onboarding: the bytes land while the user is still + * reading, and the later `load()` is a cache read. WebLLM cannot express this + * — `reload()` instantiates the wasm and needs a GPU before it fetches a + * single shard — so this is ours. See `prefetch.js` for the URL-derivation + * risk and the oracle that closes it. + * + * Needs no WebGPU at all, which is the other half of the point: an app can + * warm the cache on a machine it has not yet decided can run the model. + * + * @param {string} modelId + * @param {{signal?: AbortSignal, onProgress?: Function}} [opts] + */ + prefetch(modelId: string, { signal, onProgress }?: { + signal?: AbortSignal; + onProgress?: Function; + }): Promise<{ + modelId: string; + files: number; + bytes: number; + alreadyCached: boolean; + }>; + /** + * Free a model's bytes and **keep the registry entry**, so it stays a model + * this engine knows how to get again — the distinction from + * `store.remove()`, which forgets the URL a remote model would need. + * + * Delegates for remote and prebuilt models: `deleteModelAllInfoInCache` is + * WebLLM's, covers tensors + wasm + config, and is maintained upstream. + */ + evict(modelId: any): Promise<{ + freedKeys: number; + }>; + /** + * Forget a model entirely: free its bytes **and** drop the registry entry. + * + * `evict()` first, because that is what knows how to reach the bytes for each + * source — and it has to happen before the record is deleted, since for a + * remote model the record holds the only URL those bytes can be derived from. + * Deleting the entry first would strand them in Cache Storage permanently. + */ + remove(modelId: any): Promise<{ + freedKeys: number; + }>; + /** + * Projected decode throughput for a model, in tokens per second. + * + * `basis: "measured"` once anything has actually decoded on this machine — + * the engine then knows its own achieved bandwidth and every projection is + * device-specific. Before that, `basis: "extrapolated"` from a reference + * machine, which is a starting point and says so. + * + * Decode is memory-bandwidth-bound, so this is close to the whole story: + * time per token scales with weight bytes and little else. + * + * @param {string} [modelId] defaults to the current model + */ + estimateSpeed(modelId?: string): Promise<{ + tokensPerSecond: number; + basis: "measured" | "extrapolated"; + modelBytes: number; + bytesPerSecond: number; + reference?: string; + modelId: string; + }>; + /** + * What is actually switched on right now, as opposed to what the device could + * support. + * + * The distinction matters for KV reuse in particular: `probe().kvReuse` is a + * device capability, but the decision is taken inside the engine worker, + * which is the authority. A caller debugging "why is my second turn slow" + * needs the decision, not the capability. + */ + features(): Promise<{ + kvReuse: boolean; + shaderF16: boolean; + decodeSteps: any; + multiStepDecoding: boolean; + engines: number; + maxEngines: any; + resident: string[]; + computePassBatching: number; + decode: any; + }>; + /** + * Register a model. Two shapes, one call, and the difference is only where + * the bytes come from: + * + * ```js + * // fetched from a base URL you host — an HF repo, a CDN, your own origin + * await engine.registerModel({ + * modelId: "my-model", + * model: "/models/my-model/", + * modelLib: "/models/my-model/my-model-webgpu.wasm", + * }); + * + * // read off disk. No network connection at any point, ever. + * await engine.registerModel({ modelId: "my-model", files: entries }); + * ``` + * + * Both end up as one `model_list` entry that WebLLM's own loader resolves the + * same way — the local one only differs in that its base URL is minted on + * `.invalid` and its cache is populated before the loader ever looks. + * + * That origin is the *mechanism* of the offline guarantee, not a marker of + * it: `.invalid` is reserved by RFC 6761 and can never resolve, so there is + * no code path — no bug, no eviction, no future refactor — by which a local + * model reaches the network. It fails with a DNS error instead. + * + * `files` is `{ path, file }[]`; `filesFromDataTransfer` and + * `filesFromInput` build it from a drop event or a directory picker. + */ + registerModel(spec: any): Promise; + /** + * Bring a model up, whatever form you have it in. + * + * One entry point for all three routes, because from a caller's side "load a + * model" is one intention and having to know which of `load`, + * `registerModel` and `ingestModelFolder` to reach for is a decision the + * library can make for them: + * + * ```js + * load("Llama-3.2-1B-Instruct-q4f16_1-MLC") // prebuilt or registered id + * load("https://huggingface.co/mlc-ai/Foo", { modelLib }) // a URL you host + * load({ model, modelLib }) // the same, explicit + * load({ files }) | load(fileList) | load(dataTransfer) // a folder, no network + * ``` + * + * `registerModel` and `ingestModelFolder` remain, unchanged, as the low-level + * primitives — this composes them rather than replacing them. + * + * **A URL always needs `modelLib`.** It is not guessed; see `sources.js` for + * the measurement behind that. **`defer: true`** registers the source and + * stops there, returning the record instead of the state — the manager's + * drop-now-load-later flow. + * + * Additive residency: a model already resident stays resident, so switching + * back to it costs nothing. That is only safe while the weights fit, so + * `keepResident: false` (the default) unloads whatever else is up first — + * the old single-model behaviour, and the safe one on a 16 GB machine. + * Pass `keepResident: true` to hold both, having checked the budget yourself + * with `canRun()`. + * + * @param {string | object} src an id, a URL, `{model, modelLib}`, or a folder + * @param {{keepResident?: boolean, signal?: AbortSignal, defer?: boolean, + * id?: string, modelLib?: string, modelType?: string, contextWindow?: number, + * vramRequiredMB?: number, onProgress?: Function}} [opts] + * @returns {Promise} the engine state, or the registry record when `defer` + */ + load(src: string | object, opts?: { + keepResident?: boolean; + signal?: AbortSignal; + defer?: boolean; + id?: string; + modelLib?: string; + modelType?: string; + contextWindow?: number; + vramRequiredMB?: number; + onProgress?: Function; + }): Promise; + /** + * Let a model go, at one of two depths. + * + * ```js + * unload() // the current model's VRAM; cached bytes stay + * unload(id) // that model's VRAM + * unload(id, "cache") // and delete its cached bytes, keeping the registry entry + * ``` + * + * At `"vram"` the bytes stay on disk, so loading it again costs no network — + * that is what makes switching back cheap, and the difference between this + * and `remove()`. + * + * **A bare `unload()` frees only the current model**, not every resident one. + * `unloadAll()` is the explicit form for that: freeing everything is the more + * destructive of the two readings and should have to be asked for by name. + * + * @param {string} [modelId] defaults to the current model. Omit both this and + * any resident model to no-op. + * @param {"vram"|"cache"} [level] + */ + unload(modelId?: string, level?: "vram" | "cache"): Promise<{ + status: string; + modelId: any; + progress: any; + error: any; + pool: { + size: number; + busy: number; + queued: number; + }; + /** Model ids with a live pool. `modelId` is whichever of them is current. */ + resident: any[]; + /** Latest decode probe from an engine worker; see multistep.js. */ + decode: any; + }>; + /** Unload every resident model. */ + unloadAll(): Promise<{ + status: string; + modelId: any; + progress: any; + error: any; + pool: { + size: number; + busy: number; + queued: number; + }; + /** Model ids with a live pool. `modelId` is whichever of them is current. */ + resident: any[]; + /** Latest decode probe from an engine worker; see multistep.js. */ + decode: any; + }>; + /** + * One completion. + * + * Named `complete` rather than `chat` so `engine.chat.completions.create()` + * — the WebLLM-shaped facade, Phase 2 — can take that name without a rename. + * + * @param {CompletionRequest} payload + * @param {(delta: string) => void} [onChunk] called per streamed text delta + * @returns {Promise} + */ + complete(payload: CompletionRequest, onChunk?: (delta: string) => void): Promise; + /** + * `complete()`, but the callback receives WebLLM's chunk verbatim. + * + * Exists so the `chat.completions.create()` facade can pass chunks straight + * through instead of rebuilding an envelope — which is what dropped + * `tool_calls`, flattened `logprobs` and restamped `created`. + * + * @param {CompletionRequest} payload + * @param {(chunk: object) => void} [onRawChunk] + * @returns {Promise}>} + */ + completeRaw(payload: CompletionRequest, onRawChunk?: (chunk: object) => void): Promise; + }>; + /** + * One question, one answer, nothing kept. + * + * ```js + * const answer = await engine.ask("Summarise this in one line:\n" + doc); + * ``` + * + * @param {string | Array} input + * @param {object} [opts] anything `complete()` takes, plus `onDelta` to stream + * @returns {Promise} + */ + ask(input: string | Array, opts?: object): Promise; + /** + * A multi-turn conversation that keeps its own history. + * + * ```js + * const chat = engine.conversation({ system: "You are terse." }); + * await chat.say("hello"); + * await chat.say("and again?"); // remembers + * ``` + * + * @param {object} [opts] `system`, `keep`, plus `complete()` defaults + */ + conversation(opts?: object): { + readonly messages: any[]; + readonly length: number; + say(content: string, onDelta?: (delta: string) => void): Promise<{ + text: string; + finishReason: "length" | "stop" | "abort"; + }>; + reset(): /*elided*/ any; + restore(messages: any): /*elided*/ any; + }; + /** + * Ghost text, with the debounce/supersede/drop-if-stale discipline built in + * and the prompt left to you. + * + * ```js + * const ghost = engine.ghostText({ prompt: (before) => `Continue:\n${before}` }); + * editor.on("input", async () => { + * const hint = await ghost.suggest(editor.textBefore()); + * if (hint !== null) render(hint); // null means a newer keystroke won + * }); + * editor.on("blur", () => ghost.cancel()); + * ``` + * + * @param {object} opts must include `prompt` + */ + ghostText(opts: object): { + suggest(context: any): Promise; + cancel(): number; + }; + /** + * Embed text into vectors, through the same scheduler as everything else. + * + * ```js + * const [vector] = await engine.embed("a sentence", { modelId: EMBED_MODEL }); + * const vectors = await engine.embed(["one", "two"], { modelId: EMBED_MODEL }); + * ``` + * + * **Needs an embedding model**, not a chat model — `snowflake-arctic-embed-*` + * in WebLLM's prebuilt list, from 239 MB. They are separate models, so this + * usually names `modelId` explicitly and holds it resident alongside a chat + * model with `load(id, { keepResident: true })`. + * + * Returns bare vectors because that is what a caller does arithmetic on; the + * OpenAI envelope is available as `embedRaw()` for anyone porting code that + * expects `data[].embedding`. + * + * **A running embedding cannot be interrupted.** Cancellation and preemption + * work by making a decode loop break out; one forward pass has no loop, so a + * `cancel()` that lands after the job starts marks it cancelled but does not + * stop it. Queued embeddings supersede and cancel normally. This is tolerable + * because an embedding is milliseconds where a completion is seconds — but it + * is a weaker guarantee than `complete()` gives, so it is stated rather than + * discovered. + * + * @param {string | string[]} input + * @param {{modelId?: string, task?: string, session?: string, + * priority?: string, preemptible?: boolean, id?: string}} [opts] + * @returns {Promise} one vector per input, in order + */ + embed(input: string | string[], opts?: { + modelId?: string; + task?: string; + session?: string; + priority?: string; + preemptible?: boolean; + id?: string; + }): Promise; + /** `embed()`, returning WebLLM's OpenAI-shaped envelope untouched. */ + embedRaw(input: any, opts?: {}): Promise<{ + data: any; + usage: any; + }>; + /** + * Independent prompts, fanned across the pool. This is the only way to beat + * the ~10 tok/s single-stream ceiling, so anything embarrassingly parallel + * (translating a page, labelling a list) should arrive here rather than as a + * loop of `complete` calls. + * + * @param {CompletionRequest & {requests: Array>}} payload + * @param {(item: BatchItem) => void} [onItem] called as each item lands + * @returns {Promise>} + */ + batch(payload: CompletionRequest & { + requests: Array>; + }, onItem?: (item: BatchItem) => void): Promise>; + /** + * Cancels by job id or by session key. + * @param {string} idOrSession + * @returns {number} how many jobs it stopped + */ + cancel(idOrSession: string): number; + /** + * Applies a runtime knob to the running pool and persists it as the default. + * + * `decodeSteps` is the multi-step decode width (AI.md, "Multi-step decoding"). + * It takes effect on the next burst — no reload — which is what makes sweeping + * it to find this machine's tick boundary cheap. + */ + configure(patch: any): Promise<{ + settings: { + decodeSteps: number; + engineCount: number; + }; + engines: number; + }>; + #private; +} +/** + * The OpenAI generation fields WebLLM already speaks, plus the scheduling + * fields that are what this engine adds over calling WebLLM directly. + */ +export type CompletionRequest = { + messages: Array<{ + role: string; + content: string; + }>; + /** + * load this model first if it is not the live one + */ + modelId?: string; + /** + * job id; also what `cancel(id)` takes + */ + id?: string; + temperature?: number; + max_tokens?: number; + response_format?: object; + extra_body?: object; + /** + * the unit that owns an engine; a whole batch shares one + */ + task?: string; + /** + * a later job with this key supersedes the earlier one + */ + session?: string; + priority?: "interactive" | "normal" | "background"; + /** + * may be interrupted by an `interactive` job + */ + preemptible?: boolean; +}; +export type CompletionResult = { + text: string; + usage?: object; + /** + * WebLLM's own values + */ + finishReason?: "stop" | "length" | "abort"; + /** + * superseded or explicitly cancelled + */ + cancelled?: true; + /** + * an `interactive` job took the slot; `text` is partial + */ + preempted?: true; +}; +export type BatchItem = CompletionRequest & { + index: number; + engineIndex: number; + startedAt: number; + finishedAt: number; + error?: string; +}; +import { ModelStore } from "./model-store.js"; diff --git a/types/engine/environment.d.ts b/types/engine/environment.d.ts new file mode 100644 index 0000000000000000000000000000000000000000..80d8d856db3efd0e973debd0fd526c417799c6b7 --- /dev/null +++ b/types/engine/environment.d.ts @@ -0,0 +1,43 @@ +/** + * Builds the callable `engine.environment` — a function with `.measure()` on it. + * + * A function rather than an object because the common case is asking for the + * whole report, and `environment()` should be the short spelling of that. + * + * @param {import("./engine.js").ScheduledEngine} engine + */ +export function environmentFacade(engine: import("./engine.js").ScheduledEngine): { + (opts?: {}): Promise<{ + lines: { + id: any; + severity: any; + affects: any; + cause: any; + fix: any; + operable: boolean; + }[]; + speed?: any; + scope: string; + ok: boolean; + severity: any; + device: any; + }>; + measure({ tokens }?: { + tokens?: number; + }): Promise<{ + tokens: number; + tokensPerSecond: number; + basis: "measured" | "extrapolated"; + modelBytes: number; + bytesPerSecond: number; + reference?: string; + modelId: string; + }>; +}; +export namespace SEVERITY { + let BLOCKED: string; + let DEGRADED: string; + let TUNE: string; + let INFO: string; + let OK: string; +} diff --git a/types/engine/errors.d.ts b/types/engine/errors.d.ts new file mode 100644 index 0000000000000000000000000000000000000000..8a593ce21e5849429f85f29b7d75bf56861e5fbf --- /dev/null +++ b/types/engine/errors.d.ts @@ -0,0 +1,37 @@ +/** + * Normalise anything thrown into an EngineError. + * + * Errors from WebLLM, the GPU stack and the structured-clone boundary arrive as + * plain Errors, DOMExceptions or strings, and a caller should not have to care + * which. Anything unrecognised becomes `GENERATION_FAILED` rather than being + * given a more specific code it has not earned. + */ +export function asEngineError(err: any, fallback?: string): EngineError; +export namespace ERROR { + let NO_WEBGPU: string; + let NO_MODEL: string; + let UNKNOWN_MODEL: string; + let CACHE_INCOMPLETE: string; + let INVALID_MODEL_FOLDER: string; + let BAD_REQUEST: string; + let ABORTED: string; + let GENERATION_FAILED: string; + let PACKAGE_INCOMPLETE: string; +} +export class EngineError extends Error { + /** + * @param {string} code one of ERROR + * @param {string} message human-readable, safe to show a user + * @param {object} [detail] structured context, so callers need not parse `message` + */ + constructor(code: string, message: string, detail?: object); + code: string; + detail: any; + /** Wire form. `error` stays a plain string so existing callers keep working. */ + toJSON(): { + detail?: any; + code: string; + message: string; + }; +} +export function isEngineError(err: any, code: any): boolean; diff --git a/types/engine/index.d.ts b/types/engine/index.d.ts new file mode 100644 index 0000000000000000000000000000000000000000..4d80de0130d2d3073c10788ffe4b549ff716be6e --- /dev/null +++ b/types/engine/index.d.ts @@ -0,0 +1,12 @@ +export { ScheduledEngine } from "./engine.js"; +export { CreateScheduledEngine } from "./create.js"; +export { EnginePool } from "./pool.js"; +export { SEVERITY } from "./environment.js"; +export { canRun, probeDevice, projectSpeed, rankModels, REFERENCE_DECODE_BYTES_PER_SECOND } from "./device.js"; +export { ERROR, EngineError, asEngineError, isEngineError } from "./errors.js"; +export { ModelStore, MODEL_TYPE, SOURCE, isInjected, DEFAULT_SETTINGS, CACHE_CONFIG, CACHE_MODEL, CACHE_WASM, baseUrlFor, groupKeysByScope, toAppConfig, formatBytes } from "./model-store.js"; +export { filesFromDataTransfer, filesFromInput, ingestModelFolder } from "./ingest.js"; +export { prefetchModel, resolveModelUrl } from "./prefetch.js"; +export { ask, conversation, ghostText } from "./recipes.js"; +export { DEFAULT_DECODE_STEPS, MAX_DECODE_STEPS, burstSize, clampSteps, installMultiStepDecoding } from "./multistep.js"; +export { ENGINE_STATE, PRIORITY, PRIORITY_ORDER, UNLOAD_LEVEL, WORKER_CONFIGURE } from "./constants.js"; diff --git a/types/engine/ingest.d.ts b/types/engine/ingest.d.ts new file mode 100644 index 0000000000000000000000000000000000000000..ddccfcbe5d48117abdf52dcfd77e5504547c82a3 --- /dev/null +++ b/types/engine/ingest.d.ts @@ -0,0 +1,50 @@ +/** + * Walk a DataTransfer from a drop event into flat `{ path, file }` entries. + * Uses the entries API so dropping a *folder* works, not just a file selection. + */ +export function filesFromDataTransfer(dataTransfer: any): Promise; +/** + * Turns `` output into the same `{ path, file }` shape. + * + * `Array.from`, not spread: a real `FileList` is iterable, but plenty of things + * that behave like one are only array-like, and spreading those fails with + * "fileList is not iterable" — an error that names none of the three places it + * could have come from. Array.from accepts both. + */ +export function filesFromInput(fileList: any): { + path: any; + file: any; +}[]; +/** + * @param {Array<{path: string, file: File}>} entries + * @param {object} opts + * @param {import("./model-store.js").ModelStore} opts.store where the registry entry lands + * @param {string} [opts.modelId] overrides the id inferred from the folder name + * @param {"llm"|"embedding"|"vlm"} [opts.modelType] declare a vision model, or WebLLM + * rejects every image sent to it + * @param {(p: {phase: string, done: number, total: number, label: string}) => void} [opts.onProgress] + * @returns {Promise} the saved registry record + */ +export function ingestModelFolder(entries: Array<{ + path: string; + file: File; +}>, { store, modelId, modelType, onProgress }?: { + store: import("./model-store.js").ModelStore; + modelId?: string; + modelType?: "llm" | "embedding" | "vlm"; + onProgress?: (p: { + phase: string; + done: number; + total: number; + label: string; + }) => void; +}): Promise; +/** WebLLM asks for this name; older MLC exports ship `ndarray-cache.json`. */ +export const TENSOR_MANIFEST: "tensor-cache.json"; +export const LEGACY_TENSOR_MANIFEST: "ndarray-cache.json"; +export const CHAT_CONFIG: "mlc-chat-config.json"; +export namespace CONTENT_TYPES { + let json: string; + let wasm: string; + let bin: string; +} diff --git a/types/engine/model-store.d.ts b/types/engine/model-store.d.ts new file mode 100644 index 0000000000000000000000000000000000000000..e12867ebee8082d10a038875f9aa5b20953ec12f --- /dev/null +++ b/types/engine/model-store.d.ts @@ -0,0 +1,150 @@ +export function baseUrlFor(modelId: any): string; +/** Every cache key a record claims, keyed by cache scope. */ +export function groupKeysByScope(record: any): { + "webllm/config": any; + "webllm/model": any; + "webllm/wasm": any; +}; +/** Accepts `"vlm"`, `MODEL_TYPE.vlm`, or nothing. */ +export function toModelType(value: any): any; +/** + * Shape WebLLM's `appConfig` from the registry, optionally over its own + * prebuilt list. + * + * Registered records win on a model_id collision, so a developer can shadow a + * prebuilt entry — point `Llama-3.2-1B-Instruct-q4f16_1-MLC` at their own + * mirror, say — without renaming it and breaking their callers. + * + * @param {Array} models registered records + * @param {{model_list: Array} | null} [prebuilt] WebLLM's `prebuiltAppConfig` + */ +export function toAppConfig(models: Array, prebuilt?: { + model_list: Array; +} | null): { + model_list: any[]; + useIndexedDBCache: boolean; +}; +export function formatBytes(bytes: any): string; +export const CACHE_CONFIG: "webllm/config"; +export const CACHE_MODEL: "webllm/model"; +export const CACHE_WASM: "webllm/wasm"; +export namespace DEFAULT_SETTINGS { + let allowedExternalIds: any[]; + let engineCount: number; + let decodeSteps: number; + let temperature: number; + let maxTokens: number; + let systemPrompt: string; +} +export namespace MODEL_TYPE { + let llm: number; + let embedding: number; + let vlm: number; +} +export namespace SOURCE { + let PREBUILT: string; + let REMOTE: string; + let INJECTED: string; +} +export function isInjected(record: any): boolean; +/** + * @typedef {object} StorageAdapter + * @property {(key: string | string[]) => Promise} get + * @property {(items: object) => Promise} set + * + * Two methods, deliberately the exact shape of `browser.storage.local`, so the + * WebExtension adapter is a passthrough rather than a translation layer. See + * `src/adapters/` for the three implementations. + */ +export class ModelStore { + /** @param {StorageAdapter} storage */ + constructor(storage: StorageAdapter); + /** @returns {Promise>} registered model records, newest first. */ + list(): Promise>; + get(modelId: any): Promise; + /** + * Register weights the engine should fetch rather than find already cached. + * + * This is the whole "developer configures a read path" surface: `model` is a + * base URL and nothing more, so an HF repo, your own CDN, a folder served off + * your own origin and a localhost dev server are all the same call. + * + * ```js + * await store.registerModel({ + * modelId: "Qwen3.8-2B-q4f16_1-MLC", + * model: "/models/Qwen3.8-2B-q4f16_1-MLC/", + * modelLib: "/models/Qwen3.8-2B-q4f16_1-MLC/Qwen3.8-2B-q4f16_1-webgpu.wasm", + * }); + * ``` + * + * No validation of the URLs happens here — there is nothing to validate + * without fetching, and WebLLM's loader reports a bad base URL far better + * than a HEAD request would. Contrast `ingestModelFolder`, which validates + * exhaustively because it is about to write a gigabyte. + */ + registerModel({ modelId, model, modelLib, contextWindow, vramRequiredMB, modelType }: { + modelId: any; + model: any; + modelLib: any; + contextWindow: any; + vramRequiredMB: any; + modelType: any; + }): Promise; + save(record: any): Promise; + /** + * Frees an **injected** model's bytes; keeps the registry entry. + * + * Scoped to injected models deliberately. We wrote those artifacts, so we + * hold the only manifest of them — which also makes this stronger than + * WebLLM's equivalent *for this case*: `deleteModelAllInfoInCache` reads + * `tensor-cache.json` to enumerate shards, so it cannot clean up after that + * file has itself been evicted. An explicit key list has no such failure. + * + * Remote and prebuilt models are WebLLM's to delete — it fetched them and + * derives their keys exactly as its loader did. `ScheduledEngine.evict()` + * routes to whichever owns the model. + * + * @returns {Promise<{freedKeys: number}>} + */ + evictInjected(modelId: any): Promise<{ + freedKeys: number; + }>; + /** + * The registry-only primitive: drops the entry, plus the cache keys an + * **injected** record lists. + * + * A remote record claims no keys, so this frees nothing for it — WebLLM + * downloaded those bytes and derives their names. Use + * `ScheduledEngine.remove()` for a full teardown; calling this directly on a + * remote model leaks its shards, because deleting the record also destroys + * the only record of where they came from. + */ + remove(modelId: any): Promise; + /** + * Confirms the caches still hold everything the record promised. + * + * Storage eviction is silent, so this is what stands between a stale registry + * entry and WebLLM trying to fetch `local-model.invalid` mid-load. Inside a + * WebExtension with `unlimitedStorage` this was a defensive check; on an + * ordinary page origin, where a multi-GB model is evictable unless + * `navigator.storage.persist()` was granted, it is a core mechanism. + * + * @returns {Promise<{ok: boolean, missing: string[]}>} + */ + verify(record: any): Promise<{ + ok: boolean; + missing: string[]; + }>; + getSettings(): Promise; + setSettings(patch: any): Promise; + #private; +} +export type StorageAdapter = { + get: (key: string | string[]) => Promise; + /** + * Two methods, deliberately the exact shape of `browser.storage.local`, so the + * WebExtension adapter is a passthrough rather than a translation layer. See + * `src/adapters/` for the three implementations. + */ + set: (items: object) => Promise; +}; diff --git a/types/engine/multistep.d.ts b/types/engine/multistep.d.ts new file mode 100644 index 0000000000000000000000000000000000000000..1c70b2f40b851d3e89aebfbe4577c2e6dcb15c31 --- /dev/null +++ b/types/engine/multistep.d.ts @@ -0,0 +1,106 @@ +/** + * What this pipeline is missing, as sentences a reader can act on. + * Empty means a burst is safe to run. + */ +export function missingPipelineMembers(pipeline: any): string[]; +/** + * Replaces `engine.decode` with a burst-and-drain version. + * + * `decode` stays a one-token call — the caller's loop still checks + * `pipeline.stopped()` between tokens and still emits one chunk per token — but + * only one call in K actually touches the GPU. The rest drain a buffer. + * + * @param {object} engine an MLCEngine; in this project the one inside the worker + * @param {object} [options] + * @param {number} [options.steps] forward steps per sync; 1 disables the path + * @param {(info: {steps: number, tokens: number, ms: number}) => void} [options.onBurst] + * @param {(info: {missing: string[]}) => void} [options.onFallback] fired once + * per pipeline that fails the contract, before it is routed to stock decoding + * @returns {{setSteps: (n: number) => void, readonly steps: number, + * readonly fallbacks: number}} + */ +export function installMultiStepDecoding(engine: object, { steps, onBurst, onFallback }?: { + steps?: number; + onBurst?: (info: { + steps: number; + tokens: number; + ms: number; + }) => void; + onFallback?: (info: { + missing: string[]; + }) => void; +}): { + setSteps: (n: number) => void; + readonly steps: number; + readonly fallbacks: number; +}; +/** + * How many steps may run before the next stop condition *has* to be checked. + * + * `max_tokens` and the context window are countable, so they are clamped rather + * than overshot — which leaves stop tokens as the only reason a burst is ever + * rewound. Returns 1 when multi-step cannot be used at all, routing the caller + * to stock single-step decoding. + */ +export function burstSize(pipeline: any, genConfig: any, steps: any): number; +/** + * Multi-step decoding: N forward steps per GPU->CPU sync. + * + * Why this exists: decode here is not compute-bound, it is *sync*-bound. Firefox + * resolves `onSubmittedWorkDone()` / `mapAsync()` only on a 100 ms poll tick + * (AI.md, "The 10 tok/s ceiling"), and stock WebLLM needs exactly one sync per + * token — it reads the sampled token id back to JS before it can build the next + * step's input. One token per tick = 9.6 tok/s, of which ~7 ms is real compute. + * + * The fix is the one vLLM ships as `--num-scheduler-steps`: run K steps before + * paying the per-batch cost once. What makes it possible here without touching + * the compiled model is that WebLLM's sampling path is *already* on the GPU — + * `softmax_with_temperature`, `argsort_probs` and `sample_with_top_p` hand back + * an int32[1] device tensor, and `Tensor.copyFrom(Tensor)` is a device-to-device + * copy. So the sampled id feeds straight back into `embed` without ever becoming + * a JS number: + * + * embed -> decode -> penalties -> softmax -> argsort -> sample -> embed -> ... + * + * Each step stages its id into its own CPU tensor, and the burst ends with + * **one** `device.sync()`. tvmjs queues GPU->CPU copies into `pendingGPUToCPUCopy` + * and only awaits them in `sync()`, so K readbacks still cost one tick. + * + * Two things follow from the 100 ms grid, and they are why `steps` is a dial: + * + * - The win is quantized, not linear. A burst costs `ceil(K * perStepMs / 100)` + * ticks, so throughput is a sawtooth and the good values of K are the ones + * landing just under a boundary. On a 0.8B at ~7.3 ms/step that is K=13 + * (~130 tok/s); K=14 already spills into a second tick and halves it. + * - The best K shrinks as the model grows, because `perStepMs` grows. A model + * at 25 ms/step wants K=4, not K=15. + * + * Cost of the trick: the sampler cannot see its own output mid-burst. Repetition + * and presence/frequency penalties use the token history as it stood when the + * burst started, and stop conditions are only checked after the readback, so a + * burst can overshoot a stop token and must then be rewound. Both are the same + * trade vLLM makes. Anything needing per-token CPU feedback (grammar-constrained + * JSON, logprobs, a logit processor) falls back to single-step, where behaviour + * is identical to stock WebLLM. + * + * The other cost is that all of this drives ~30 undocumented tvmjs internals. A + * WebLLM upgrade that renames one does not break generation — it turns the fast + * path off and takes the throughput with it, silently. `PIPELINE_CONTRACT` below + * is that surface written down and checked against the live pipeline before the + * first burst, so the failure announces itself instead of being measured months + * later. + */ +/** vLLM's documented sweet spot, and the value this extension ships. */ +export const DEFAULT_DECODE_STEPS: 15; +/** + * Above this, the lookahead thrown away at a stop token outweighs the tick it + * saves, and the transient logits/argsort buffers stop being free. + */ +export const MAX_DECODE_STEPS: 32; +export function clampSteps(n: any): number; +export namespace PIPELINE_CONTRACT { + let calls: string[]; + let numbers: string[]; + let reads: string[]; + let optional: string[]; +} diff --git a/types/engine/pool.d.ts b/types/engine/pool.d.ts new file mode 100644 index 0000000000000000000000000000000000000000..17e70862fdb434a5c09fefc334b7cc2d66c719c8 --- /dev/null +++ b/types/engine/pool.d.ts @@ -0,0 +1,68 @@ +export class EnginePool { + constructor({ size, createEngine, onStateChange }: { + size?: number; + createEngine: any; + onStateChange?: () => void; + }); + /** Engines that exist right now, which is not the same as the cap. */ + get size(): number; + get maxSize(): number; + get loaded(): boolean; + /** + * Brings up the first engine, and only the first. + * + * The pool used to build every engine here. It no longer does, because an + * engine is a full copy of the weights - measured ~1.6 GB steady state for + * the 0.8B, ~2.4 GB for a 2B - and staging a second one costs that much host + * memory before the GPU ever sees it. On a 16 GB machine that is the + * difference between working and swapping, and it was being paid up front + * whether or not two tasks ever ran at once. `#grow` earns the rest. + */ + load(onProgress?: () => void): Promise; + unload(): Promise; + /** + * @param {object} spec + * @param {object} spec.params passed straight to `engine.chat.completions.create` + * @param {string} [spec.task] the unit that owns an engine; a whole batch shares one + * @param {string} [spec.session] later jobs with this session supersede earlier ones + * @param {string} [spec.priority] one of PRIORITY + * @param {boolean} [spec.preemptible] may be interrupted by an interactive job + * @param {(chunk: object) => void} [spec.onChunk] WebLLM's chunk, verbatim + * @returns {Promise<{text: string, usage?: object, cancelled?: boolean, preempted?: boolean}>} + */ + submit(spec: { + params: object; + task?: string; + session?: string; + priority?: string; + preemptible?: boolean; + onChunk?: (chunk: object) => void; + }): Promise<{ + text: string; + usage?: object; + cancelled?: boolean; + preempted?: boolean; + }>; + /** Cancels by job id or by session key. Returns how many jobs it stopped. */ + cancel(idOrSession: any): number; + /** + * Pushes a runtime setting to every engine that accepts one. + * + * Separate from `load()` because the settings this carries — `decodeSteps` so + * far — are per-generation knobs, not per-model ones: retuning them must not + * cost a reload of the weights. + */ + configure(patch: any): number; + status(): { + size: number; + maxSize: number; + growing: boolean; + growthBlocked: any; + busy: number; + queued: number; + queuedByPriority: { + [k: string]: number; + }; + }; + #private; +} diff --git a/types/engine/prefetch.d.ts b/types/engine/prefetch.d.ts new file mode 100644 index 0000000000000000000000000000000000000000..3484c01d4dafac5a7e208874a81af759fc8d716d --- /dev/null +++ b/types/engine/prefetch.d.ts @@ -0,0 +1,35 @@ +/** + * WebLLM's `cleanModelUrl`, mirrored deliberately. + * + * Kept character-for-character with the bundle's version (see the contract + * test) because the whole point is to produce the same cache keys its loader + * will look for. + */ +export function resolveModelUrl(modelUrl: any): string; +/** + * @param {object} opts + * @param {string} opts.modelId + * @param {object} opts.record the merged app-config entry: `model`, `model_lib` + * @param {(p: {phase: string, done: number, total: number, label: string}) => void} [opts.onProgress] + * @param {AbortSignal} [opts.signal] + * @param {(url: string, init?: object) => Promise} [opts.fetchImpl] + * @returns {Promise<{modelId: string, files: number, bytes: number, alreadyCached: boolean}>} + */ +export function prefetchModel({ modelId, record, onProgress, signal, fetchImpl, }: { + modelId: string; + record: object; + onProgress?: (p: { + phase: string; + done: number; + total: number; + label: string; + }) => void; + signal?: AbortSignal; + fetchImpl?: (url: string, init?: object) => Promise; +}): Promise<{ + modelId: string; + files: number; + bytes: number; + alreadyCached: boolean; +}>; +export function isAlreadyLocal(record: any): boolean; diff --git a/types/engine/recipes.d.ts b/types/engine/recipes.d.ts new file mode 100644 index 0000000000000000000000000000000000000000..28b68627562b76b8424f00bc14eb35062c4e15d2 --- /dev/null +++ b/types/engine/recipes.d.ts @@ -0,0 +1,107 @@ +/** + * A single question, with nothing kept afterwards. + * + * Its own task and no session, so two `ask()`s in flight never supersede each + * other and never queue behind one another for the same engine — which is what + * would happen if this shared a session key with anything else. + * + * @param {import("./engine.js").ScheduledEngine} engine + * @param {string | Array<{role: string, content: string}>} input + * @param {object} [opts] anything `complete()` takes; `onDelta` to stream + * @returns {Promise} the reply text + */ +export function ask(engine: import("./engine.js").ScheduledEngine, input: string | Array<{ + role: string; + content: string; +}>, { onDelta, ...opts }?: object): Promise; +/** + * A multi-turn conversation that remembers its own history. + * + * Two things this gets right that hand-rolled history usually does not: + * + * **One task for the whole conversation.** Every turn carries the same `task`, + * so a conversation holds at most one engine and a long reply can never occupy + * the pool while another conversation waits. Turns within it are serialised, + * which is what a conversation means anyway. + * + * **A bounded history, by default.** There is no cross-turn KV reuse on this + * stack (AI.md: every turn re-prefills the whole history at ~5.27 ms/token), so + * an unbounded conversation gets quadratically slower and a turn near the 4096 + * limit waits ~22 s for its first token. `keep` bounds it. Set `keep: Infinity` + * to opt out, having read that sentence. + * + * @param {import("./engine.js").ScheduledEngine} engine + * @param {object} [opts] + * @param {string} [opts.system] a system message, prepended and never trimmed + * @param {number} [opts.keep] how many *exchanges* of history to carry + * @param {string} [opts.task] defaults to a fresh id + */ +export function conversation(engine: import("./engine.js").ScheduledEngine, { system, keep, task, ...defaults }?: { + system?: string; + keep?: number; + task?: string; +}): { + /** The messages as they would be sent, including the system message. */ + readonly messages: any[]; + /** Exchanges currently retained. */ + readonly length: number; + /** + * Say something and get the reply, with the exchange appended to history. + * + * Serialised against the previous turn: a conversation whose turns + * overlapped would interleave history and produce replies to the wrong + * question. + * + * @param {string} content + * @param {(delta: string) => void} [onDelta] + */ + say(content: string, onDelta?: (delta: string) => void): Promise<{ + text: string; + finishReason: "length" | "stop" | "abort"; + }>; + /** Forget the history. The system message and settings survive. */ + reset(): /*elided*/ any; + /** Drop the history in, e.g. when restoring a saved conversation. */ + restore(messages: any): /*elided*/ any; +}; +/** + * Ghost text: the scheduling discipline, with the prompt left to the caller. + * + * Every part of this exists because of a specific way ghost text goes wrong: + * + * | | | + * | --- | --- | + * | one stable `session` | a fresh id per keystroke makes every stale request still run — AI.md's first "getting these wrong" row | + * | `interactive` priority | it is the one band that may preempt work that opted in | + * | debounce | a request per keystroke queues faster than the GPU drains | + * | short `max_tokens` | ghost text is a few words; paying for more is pure latency | + * | resolves `null` when stale | so a superseded suggestion **cannot** be rendered by mistake | + * + * That last one is the difference between this and a wrapper. The engine + * already supersedes stale requests; what a caller still has to remember is not + * to paint the answer that comes back. Returning `null` removes the choice. + * + * @param {import("./engine.js").ScheduledEngine} engine + * @param {object} opts + * @param {(context: any) => string | Array} opts.prompt **required** — + * builds the messages. Never defaulted: see the module header. + * @param {number} [opts.debounceMs] + * @param {number} [opts.maxTokens] + * @param {string} [opts.session] + */ +export function ghostText(engine: import("./engine.js").ScheduledEngine, { prompt, debounceMs, maxTokens, session, ...defaults }?: { + prompt: (context: any) => string | Array; + debounceMs?: number; + maxTokens?: number; + session?: string; +}): { + /** + * Ask for a suggestion. Debounced, superseding, and `null` when stale. + * + * @param {any} context whatever `prompt` takes + * @returns {Promise} `null` if superseded or cancelled + */ + suggest(context: any): Promise; + /** On blur, or on accept. Cancels in flight and invalidates anything pending. */ + cancel(): number; +}; diff --git a/types/engine/sources.d.ts b/types/engine/sources.d.ts new file mode 100644 index 0000000000000000000000000000000000000000..d6b59b437b3af948a88808facf2ff12a38460124 --- /dev/null +++ b/types/engine/sources.d.ts @@ -0,0 +1,42 @@ +/** + * The local label for a model loaded from a URL. + * + * Safe to derive, unlike `modelLib`: an id is a key in our own registry, never + * a path anything fetches, so a wrong guess is visible immediately and costs + * nothing. `.../mlc-ai/Foo-MLC` and `.../models/foo/` both give the last real + * segment. + */ +export function idFromUrl(url: any): string; +/** + * @param {string | object} src + * @param {{id?: string, modelLib?: string}} [opts] + * @returns {{kind: SourceKind, modelId?: string, model?: string, modelLib?: string, files?: unknown}} + */ +export function classifySource(src: string | object, opts?: { + id?: string; + modelLib?: string; +}): { + kind: SourceKind; + modelId?: string; + model?: string; + modelLib?: string; + files?: unknown; +}; +/** + * Ids close enough to be a typo, so an unknown id can say "did you mean". + * + * Cheap on purpose — case-insensitive substring both ways, then a bounded edit + * distance. The list is at most a few hundred entries and this runs once, on a + * path that is already about to throw. + */ +export function nearMatches(wanted: any, available: any, limit?: number): any; +export namespace SOURCE_KIND { + let ID: string; + let REMOTE: string; + let REGISTER: string; + let FILES: string; +} +export function looksLikeUrl(s: any): any; +export function isDataTransfer(v: any): boolean; +export function isFileList(v: any): boolean; +export type SourceKind = "id" | "remote" | "register" | "files"; diff --git a/types/vite.d.ts b/types/vite.d.ts new file mode 100644 index 0000000000000000000000000000000000000000..60adb7d0d0364bd2519863d274bd1a4c6cd16e63 --- /dev/null +++ b/types/vite.d.ts @@ -0,0 +1,10 @@ +/** + * @param {{exclude?: string[]}} [opts] `exclude` adds further specifiers, for a + * host that re-exports this engine from its own package and hits the same + * pre-bundling of the worker URL. + * @returns {import("vite").Plugin} + */ +export function everythingWebGPU({ exclude }?: { + exclude?: string[]; +}): any; +export default everythingWebGPU;