Skip to content

feat: add decision heads for Laya models (apiVersion 1) - #112

Merged
leehack merged 1 commit into
mainfrom
feat/decision-head
Sep 23, 2026
Merged

leehack merged 1 commit into
mainfrom
feat/decision-head

Conversation

@leehack

@leehack leehack commented Sep 23, 2026 •

Copy link
Copy Markdown
Owner

Part of llamadart#604: Web support for llamadart's DecisionEngine (Laya decision models).

What changed

LlamaWebGpuBridge gains a decision-head API, contract apiVersion 1. It runs on the direct and worker runtimes, and every call goes through the exclusive-operation queue:

  • getDecisionCapabilities(): {apiVersion, supported, reason?} for the loaded model. A model is supported only if it is a ModernBERT encoder with CLS, SEP and MASK tokens.
  • loadDecisionHead(source, {configJson?, onProgress?}): loads from a URL or from bytes and returns the handle, hidden size, special tokens, MASK text, config text and head device.
  • runDecision(handle, sequences): returns raw marker logits and actLogits per sequence.
  • freeDecisionHead(handle): unknown handles are ignored.

src/llama_webgpu_decision.{h,cpp} ports llamadart's native decision head:

  • Encoder: each head gets a private ModernBERT encoder context (pooling NONE, n_ctx = n_batch = n_ubatch = max_len).
  • Head: its weights sit in one backend buffer, and its graph runs through ggml_backend_sched with the CPU backend last. The act MLP runs on the host with std::erf.
  • Checks at load and run match llamadart's native checks, plus a limit of no more markers than tokens:
    • architecture, special tokens and widths;
    • n_ctx_train, pooling and n_ubatch;
    • token count before llama_encode, marker range, question type and vocabulary.
  • Untrusted input: safetensors parsing is bounded (F32/F16/BF16), and JSON input is limited to 64 nesting levels.
  • Data transfer: inputs, outputs and configJson go through WASMFS files, because ccall strings are copied onto the 1 MiB stack. Staged files are removed with unlink, without an analyzePath read-back.
  • Lifetime: heads are freed before the model in free_runtime, on every model load and at shutdown. Handles are never reused. Load and run are refused while generation or TTS is active.

JS side:

  • A facade handle only works with the worker or runtime that loaded it.
  • Worker head loads re-arm their timeout on download progress. Worker runs allow 10 min plus 60 s per sequence.
  • A worker failure during a run falls back to the main thread and invalidates the handles.

Types in llama_webgpu_bridge.d.ts; docs in docs/api.md, README, AGENTS and CONTRIBUTING. New checks:

  • scripts/verify_decision_api.py (static contract);
  • scripts/decision_bridge_contract_test.mjs (stub-core JS contract, part of check:js);
  • scripts/decision_browser_smoke.py (real-model parity smoke; local-only, not run in CI).

All are wired into CI, ci_scope.py, the candidate and publish workflows, and verify_ci_reliability.py.

The release manifest does not advertise a decision capability. Adding one would make strict readback reject every manifest published before it (a test now pins this), and no qualification gate runs decision heads. Consumers probe with getDecisionCapabilities().

Verification

Toolchain: Homebrew emcc 6.0.10-git. The emsdk.version pin (6.0.8) is unchanged; the version check was skipped for local builds only. Builds, all clean with no warnings in bridge sources:

  • llama.cpp b29c606 (v0.4.1): wasm32 and memory64.
  • Pinned v0.2.0: wasm32 and memory64.
  • v0.4.0: wasm32.

Checks:

  • npm run check:js: passes, and the regenerated js/* is byte-identical.
  • verify_decision_api.py, verify_text_to_speech_api.py, verify_state_persistence_api.py, verify_ci_reliability.py, mtmd_compat_contract_test.py, wasm64_runtime_patch_contract_test.py and release_qualification_test.py: all pass.
  • python3 -m unittest discover -s scripts -p '*_test.py': 416 tests OK.
  • JS mutants: 16 of 16 killed. Native ASan harness: every deep-JSON case is rejected without overflow.

Smokes:

  • State-persistence smoke: passes in direct and worker mode on v0.4.1 and v0.2.0.
  • Speech-to-text smoke (Qwen3-ASR): passes in all 4 modes.
  • Multimodal smoke: not run as specified, because the pinned projector is not on disk. LFM2-VL-450M fails the same way on base and on this branch.

decision_browser_smoke.py against the Laya 0.3.5 PyTorch reference (24 rows):

  • Modes: wasm32 and wasm64, direct and worker.
  • Hostile-input stages: 513 tokens, deep header and deep config, excess markers, a large config, invalid bytes, freed handles.
  • The smoke fails on the pre-fix build.
  • Results: no choice or act decision changed, and all modes gave identical diffs.
Encoder + head Backend Worst logit diff Worst prob diff Act relative diff
Q8_0 + laya-head.safetensors WebGPU 0.1636 0.0180 0.0111
F16 + laya-head.safetensors WebGPU 0.0169 0.0027 0.0012
Q8_0 + official model.safetensors + config WebGPU 0.1636 0.0180 0.0111
Q8_0 + laya-head.safetensors CPU 0.2326 0.0552 0.0208
F16 + laya-head.safetensors CPU 0.0149 0.0016 0.0012

llamadart native Metal gives 0.0118 (F16) and 0.1642 (Q8_0) on the same rows. The head gives the same numbers on WebGPU and CPU, so the differences come from the encoder.

Lifetime:

  • 45 load/run/free cycles in wasm32 on GPU and CPU: no OOM, no leftover files.
  • A non-encoder model is reported unsupported.
  • Two heads can be live at once.
  • dispose() with a GPU head loaded is clean.

Review:

  • A build/review/fix pass fixed 8 findings. The blocker among them was the manifest capability.
  • An independent pre-commit review then found nothing blocking. Its non-blocking notes on the file read-back, the reload wording and the worker clone are fixed in this commit.

Out of scope

  • Cancelling a run.
  • Found along the way, all pre-existing:
    • #109 worker timeout options are ignored;
    • #110 a second model or projector load in one runtime fails with "FS error";
    • #111 large string arguments overflow the wasm stack.
  • Asset publication and the llamadart pin bump. The llamadart wiring, still in progress, probes for these methods, so older assets report decision models as unsupported.

Adds getDecisionCapabilities, loadDecisionHead, runDecision and
freeDecisionHead on the direct and worker runtimes. Each head gets a private
ModernBERT encoder context and runs its graph through ggml_backend_sched with
the CPU backend last; the act MLP runs on the host. Inputs, outputs and config
text pass through WASMFS files. Heads are freed before the model.
@leehack
leehack marked this pull request as ready for review September 23, 2026 18:20
@leehack
leehack merged commit 64ba825 into main Sep 23, 2026
5 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant