mirror of https://github.com/Helldez/BigMoeOnEdge.git synced 2026-10-03 03:25:42 +00:00

feat(engine): self-speculative decoding — the model's own MTP head, or n-gram lookup (#134 )

* feat(engine): MTP self-speculative decoding for Qwen3.5/3.6 (proposal)

Qwen3.5/3.6 ship a trained multi-token-prediction block inside the gguf. With
--mtp that head drafts --mtp-draft continuation tokens and the target verifies
all of them in one wider decode, confirming the longest prefix whose argmax
equals what the target itself would have produced. Nothing is approximated and
no weight is skipped, so the quality is the full model's — but it is NOT
byte-identical the way --overlap and --prefetch are, and must not be used in a
byte-identity gate: a verify pass evaluates 1+N positions in one batch, and a
batched matmul is not bit-identical to N single-token ones, so a near-tie can
flip. Off by default.

The prize is that a decode's dominant cost, moving the dense weights and the
routed expert slices, is paid once per group instead of once per token. The
counterweight is that the verify positions route independently, so a layer's
read set widens toward N*k wherever adjacent tokens disagree, and the draft
pass routes through the MTP block's own expert layer on top. Measured on the
desktop host (DRAM-bandwidth-bound, model streamed at ~1.4x RAM): +15.1% at
draft 3 with the host's best recipe (7.12 -> 8.19 tok/s), +29% without the
lossy drop knob, acceptance falling from 71% at draft 2 to 52% at draft 4, and
flash bytes per token rising 19.7 -> 33.7 MiB as the widening predicts. Draft 3
is the optimum here; 4 is worse than 2. On a flash-I/O-bound phone that balance
can invert, so the flag ships off pending the device A/B.

The orchestration is llama.cpp's own (common/speculative.h, public headers
only): no fork, no patch, no submodule bump. Self-speculation is one model with
two contexts over it — the target, created with n_rs_seq so a rejected tail is
rewound from a bounded snapshot rather than replayed, and a draft context
created with ctx_type = LLAMA_CONTEXT_TYPE_MTP. The engine builds the draft
context itself rather than through common_speculative_init_from_params because
the eval callback is per-context: the streamer only sees the MTP block's expert
layer if the draft context carries the same cb_eval.

The MTP block is streamed like any other layer. It sits at layer index n_layer,
contiguous with the trunk and using the same tensor naming, so the hook and the
expert source are sized n_layer + n_layer_nextn; left at n_layer its experts
stay silently mmap-resident. Two consequences that are easy to get wrong: the
capture warm-up has to run on the draft context too (the MTP graph is built
nowhere else), and prefill is fed through the driver so the draft context's KV
reaches the last prompt position.

The loop accepts BEFORE catching the draft context up, so the catch-up runs on
the accepted prefix instead of the whole verify batch. Acceptance depends only
on the target's logits, which are already in hand once the decode returns, and
the rejected tail was being computed only to be deleted a few statements later.
The resulting state is identical — the driver seeds from row
min(n_accepted, n_rows-1), the same row under either batch, and the surviving
KV is exactly the range the rollback used to carve out — while skipping
n_draft - n_accepted positions through the MTP block per group. Since that
block carries its own MoE FFN, on a streamed device those are expert reads that
no longer happen. It also removes the draft context's rollback entirely: it is
never given a tail to drop.

Requires an MTP-converted gguf (most quantisations strip the nextn tensors) and
greedy decoding; both are rejected at load with a message rather than silently
ignored, as is a n_ubatch narrower than the verify batch, which would split the
graph back into single-token passes and spend the draft for nothing.

Telemetry: an "mtp:" summary line, an mtp_batch per-token CSV column (a verify
decode's whole cost is charged to its group's first row, the rest carry zeros),
mtp_drafted / mtp_accepted / mtp_decodes in the CSV trailer and in BMOE_DONE,
and mtp / mtp_draft_max in the CSV preamble. The Android app exposes the flag
and the draft width, off by default.

Host gates pass. Validated on Qwen3.6-35B-A3B-MXFP4 with the streamed recipe:
draft 1 and draft 3 produce identical text, which is the invariant a broken
accept/rollback path would violate. Device A/B still owed.

* perf(mtp): shrink the draft context, make its cost measurable, record the device verdict

The first on-device A/B says MTP loses at every draft width, and the counters
say why. Same gguf with the flag on and off, shipping recipe (overlap, 3000 MiB
cache, pinned dense, drop 0.75), Qwen3.6-35B-A3B-Q4_K_M streamed:

off 5.82 / 6.14 tok/s 69.3 MiB/tok 69-109 majflt/tok
--mtp-draft 2 5.59 93.8 230
--mtp-draft 3 4.38 106.6 633

Speculation is working - 2.35-2.52 tokens per verify decode, 52-69% acceptance
- and still losing, because the prize does not exist in this regime.
stall_s/tok is 0.025-0.027 in every one of those runs, MTP on or off: 11-16% of
the token. This configuration is compute-bound, and what MTP amortises is weight
movement. The costs meanwhile are real and monotonic in the draft width: the read
set widens (+35%, +54% flash bytes per token), CPU per token rises (+28%, +67%),
and the draft context's memory tips the device into a fault storm.

Two things follow, and both are engine bugs rather than facts of nature.

The draft context's graph width drops from 256 to 32. Compute buffers are
reserved for the widest ubatch and the dominant term scales with
ubatch x vocabulary; on device that reservation measured 493 MiB - for a context
that evaluates ONE token per draft step and is handed at most 1 + draft_max
positions by the catch-up, with no logits asked for. Only prefill ever feeds it a
wide batch, and that is one layer, so splitting it costs very little. On this
engine memory is never free: it is the expert cache's, and the cache is what
decides whether the widened verify read set is a hit or a flash read.

And the cost of speculation is now measured instead of inferred. Drafting happens
between decodes, so it never entered wall_ms and tok/s never included it - a
speculated run could report a rate the user was not experiencing. New
mtp_draft_ms per-token column (a slice of loop_overhead_ms, not an addition),
mtp_draft_s/tok in the CSV trailer, mtp_draft_s_tok and loop_overhead_s_tok in
BMOE_DONE, and a second "mtp:" summary line printing the effective rate next to
the reported one.

Adds --mtp-p-min F, which stops drafting once the head's confidence in what it is
proposing falls below F. The draft loop already had this floor and the engine was
passing 0, so it always drafted the full width however unsure the head was - with
roughly half the drafts rejected at draft 3, that is the cheapest waste available
to cut. On a streamed device it pays twice: a draft not made is a pass through the
MTP block (which carries its own MoE FFN, so its own expert reads) that never
happens, AND one fewer independently routed position in the verify batch. Default
0, the setting the host numbers were measured at; the useful value is a property
of a device's balance between drafting cost and acceptance, so it is a knob to
measure rather than a constant to guess.

The Android app now reads the mtp_* keys it was already being sent: acceptance,
tokens per pass, and the effective rate. Before this the UI could not tell whether
speculation had run at all - only the session CSV could - which made the A/B this
commit reports impossible to run from the phone.

Neither mitigation changes the regime. The honest expectation is nearer
break-even, not a win, and the flag stays off by default.

Host gates pass. Note the noise floor: the two off runs did byte-identical work
and still differ by 5.6% in tok/s, and the runs were back-to-back without thermal
gating - the mechanism counters are the trustworthy part, not the exact deltas.

* perf(mtp): split the drafting flash cost from the widened verify batch

A speculated run streams more bytes per token for two unrelated reasons: the
MTP block carries its own MoE FFN, so every draft pass routes experts of its
own, and the verify batch widens the trunk's read set wherever adjacent
positions disagree. They need opposite fixes -- a narrower draft attacks the
first, only better agreement attacks the second -- and the route trace can
separate neither, since its framing brackets the target decode while the head
only ever runs in the draft context.

Measure the head's share directly by bracketing both drafting passes with the
expert source's byte counter, and report it as a third mtp: summary line.

Also record the branch-deletion rule in AGENTS.md: a branch list should only
show work in flight, and a rejected PR loses nothing.

* feat(engine): n-gram prompt-lookup draft source, and the measurement that closes it

The flash split added last commit said where MTP's cost actually is: at draft 3 on
the host, the head's own routing was 2.9% of the extra bytes a speculated run
streams and the widened verify batch was the other 97.1%. So a cheaper draft
producer is worth almost nothing, and the only property that could matter is one
the head does not have -- the ability to decline to draft at zero cost.

--ngram is that source. It takes the last few tokens, finds where that run occurred
before in the prompt or in what has been generated, and proposes whatever followed.
No head, no draft context, no decode, no expert read, and it works on any gguf
including the ones --mtp refuses for want of a nextn block. Below --ngram-min-match
it proposes nothing and the step falls through to a plain single-token decode.

Measured on the host, Qwen3.6-35B-A3B-MXFP4 streamed, 256 greedy tokens, cells
back-to-back with off run twice:

prose off 5.80 / 6.59 mtp3 7.32 eff ngram3 6.51 (cov 7.4%)
copy-heavy off 5.45 / 5.65 mtp3 6.43 eff ngram3 5.24 (cov 15%)

The zero-cost claim holds exactly -- mtp_draft_s/tok reads 0.0000 in every n-gram
cell, against 0.020-0.023 for the head plus the ~500 MiB of expert cache its draft
context takes. But the floor turns out to be per STEP, not per run: the 15% of steps
that did draft widened the read set to 67.2 MiB/token against 48-58 at baseline and,
at 44% acceptance, bought 1.20 tokens per decode. That is not enough to earn the
widening back, and a modest fraction of such steps sinks the run.

A --ngram-min-match sweep settles it rather than leaving it open. Raising the gate
3 -> 5 -> 8 lifts acceptance 44% -> 75% while coverage collapses 15% -> 3.4%, and
narrowing to --draft 1 reaches 82.6% -- the head's own figure on this prompt. Every
cell climbs toward baseline from BELOW and none crosses it; the best configuration
found lands on the floor. A knob whose optimum is its own disablement is not a
tuning problem. Acceptance, not drafting cost, is what pays for a widened batch, and
what a trained head buys is being right often enough to justify a batch that has
already been widened.

--ngram ships off. It is kept because it is the only speculation available on a
model with no head, because the per-step floor is real, and because the counters it
adds make the next speculation claim falsifiable.

Wiring. MtpConfig became SpecConfig with DraftSource {none, mtp, ngram}, and
--mtp-draft became --draft: the width belongs to the verify batch, not to whoever
filled it. --mtp and --ngram are rejected together rather than resolved by flag
order. In the session the gate split in two -- spec_on (wide batch, acceptance,
rollback: both sources) against mtp_on (draft context, common/speculative.h, the
catch-up: the head only) -- which is what lets the n-gram source reuse the whole
verify half while allocating nothing.

A step that drafts nothing now takes the plain path: llama_batch_get_one with a
logits row of -1, byte for byte the unspeculated decode. It used to build the wide
batch anyway. Required for --ngram, and it tightens --mtp-p-min's zero-draft steps
for free.

The matcher is pure policy over token ids with no llama.cpp at all -- not even
llama.h, since llama_token is int32_t -- so it sits on the clean side of the seam,
adds no dependency on the common layer, and is unit-tested with no model
(tests/ngram_test.cpp covers tie-breaks, clipping, self-match exclusion and the gate
boundary). Telemetry: spec= / spec_draft_max= / ngram_min_match= in the CSV
preamble, a new drafted_steps key in the trailer and BMOE_DONE, and an ngram: line
reporting coverage -- without which a delta cannot be divided by the fraction of the
run it applies to. The per-token and trailer counters keep their mtp_ names: they
always described the loop rather than a source, spec_* already means the temporal
prefetch in that trailer, and renaming would break every CSV already holding a
measurement. The Android setting became a three-way picker, migrating the old
boolean preference.

The device A/B agrees and adds a cost the host could not show. Thermally gated cells
(a 120 s settle, then a battery-temperature gate, so all six start between 35.3 and
36.4 C): prose 4.90 inside a 4.59-5.17 band, copy-heavy 3.14 against 4.43 -- a 29%
loss, worse than MTP's 18%. Major faults per token go 126 -> 1427 for a source that
allocates no draft context at all, and that is the rollback snapshots: n_rs_seq =
draft_max is asked for by ANY speculation, since rejecting a draft means rewinding the
KV, and on a hybrid attention/SSM model that snapshot is a real allocation scaling with
the context. The n-gram source escapes MTP's draft context but not the loop's own
memory, and on device that memory is the expert cache's.

The same run re-measured MTP with the thermal confound removed -- 3.64 effective
against 4.43, so the earlier device verdict was not an artefact of benching without a
cooldown gate -- and reproduced the flash split at 3.7% head against 96.3% widened
verify batch, matching the host's 2.9-3.0%.

Byte-identity gates pass; speculation stays out of them for the reason docs/mtp.md
gives.

The app's CSV configuration surface follows: the three new preamble keys get their own
glossary entries rather than falling through to the unexplained-key renderer, and the
draft source joins the short run label. A speculated run is not the same KIND of run --
under speculation a decode confirms a whole group, so its per-token rows are not even
accounted the same way -- and two compare legends differing by it must not read alike.

2026-08-02 00:09:39 +02:00

6.1 KiB

Raw Blame History

Working on BigMoeOnEdge (agent guide)

Read this before making changes. It captures the invariants that keep this project clean. It follows the AGENTS.md convention, so any coding agent picks it up; CLAUDE.md just points here.

What this is

A ports-and-adapters engine that streams MoE experts from flash so >RAM models run on device, built on top of llama.cpp's public API. The whole value proposition is that we do not fork llama.cpp. See docs/architecture.md and docs/seam.md.

Project map

core/include/bmoe/ — ports (interfaces) + config. Pure policy, no llama.cpp include.
core/src/io/ — platform_io (cross-platform O_DIRECT reads + reserve/commit/evict VM); file_reader (pooled positioned reader, per-consumer O_DIRECT — used by both the expert stream and the dense loader).
core/src/moe/ — gguf_offsets, arch_registry, expert_stream_source, router_hook; dense_weights (the non-expert weight policy: mmap / warm / anon, plus the residency sensor).
core/src/engine/runtime.cpp — composition + greedy generation loop.
cli/main.cpp — bmoe-cli; the ONLY place environment variables are read.
third_party/llama.cpp — stock upstream submodule.
tests/ — byte-identity gates. examples/android/ — the demo APK.

Build and test

git submodule update --init --recursive
scripts/build-host.sh
cd build && ctest --output-on-failure         # byte-identity gates (needs python3 + gguf)

Android CLI: pwsh scripts/build-android.ps1 (needs the NDK), then build the APK in examples/android.

Hard rules

Never patch llama.cpp in-tree. Everything goes through the public eval-callback and public gguf/model APIs. If a change seems to need a llama.cpp edit, stop and discuss — the fallback is a separate 1-commit fork branch on Helldez/llama.cpp, never an in-tree diff, and only after agreement. Upgrading llama.cpp must stay a submodule bump.
Repack stays off. The engine loads with use_mmap=true, use_extra_bufts=false. The streamer rebinds tensor->data to the native gguf layout; repacking breaks it. This is load-bearing, not a tunable.
No env vars in the library. core/ never calls getenv. Config flows through RunConfig; the CLI resolves any env overrides before building it.
No hardcoding. New architectures are recipe rows in arch_registry.cpp; expert counts, strides and offsets are discovered at runtime. No model-specific constants in the streaming path.
Gates must pass before merge. bmoe_moe_gates proves streamed == resident. If you touch the streamer, the seam, or bump the submodule, run them.
Docs and changelog ship with the change. Every PR updates CHANGELOG.md and the docs it invalidates, in the same PR — never as a later sweep. A release gets its own dated ## [X.Y.Z] - YYYY-MM-DD section; nothing accumulates under [Unreleased]. Check in particular: the README benchmark tables and model list, the docs/architecture.md layer map, docs/seam.md when the llama.cpp boundary moves, docs/telemetry.md when CSV columns or the BMOE_* protocol change, docs/roadmap.md when a listed future item ships, and examples/android/README.md when the catalog, settings or build flow change. Docs that name a file the code no longer has are worse than no docs.
Review exactly what is being published, every time. This repo is public and every push is permanent record. Before any commit, push, PR or release: run git status --short and stage by explicit path only — never git add -A / git add .; untracked files in the working tree are not yours to publish. Logs, CSVs and bench evidence get a scan for identifying data (device model codes, local paths, addresses) before landing in docs/; phrase the test device generically. After a squash-merge, verify the landed tree (git ls-tree) before pushing anything else.

Conventions

Commits: Conventional Commits (feat:, fix:, docs:, refactor:, test:, build:, ci:, chore:). Author is Helldez only — do NOT add AI co-author or session trailers to commits.
Group commits. One commit = one coherent change. Never a commit for a trivial tweak on its own — fold small fixes, doc touches and follow-ups into the change they belong to. If several small things accumulate, batch them into one commit.
Delete the branch when its PR closes — merged or rejected, local and remote (git branch -d, git push origin --delete). A branch list should only show work in flight. Nothing is lost on a rejected PR: GitHub keeps its commits reachable from the closed PR itself.
Language: all code, comments, docs, and commit messages in English.
Style: .clang-format (LLVM base, 4-space, 120 col). CI checks with clang-format 18; match that version locally (pip install clang-format==18.*) or formatting that looks clean can still fail the check.
Comments explain why / invariants, not what.
No milestone codenames (M0…Mn) in docs — describe capabilities thematically.

Releases

Release APKs come from CI, never from a local build. The release-apk workflow runs when a release is published: clean checkout of the tag, NDK build, signed with the stable key from repository secrets, assets attached to the release. Do not hand-upload an APK.
Every released feature bumps the app version: versionCode + versionName in the Android app's Gradle config, in the same PR as the change being released, matching the tag.
Release title is the bare version — vX.Y.Z, no description after it.
Validate on device before releasing. The host gates prove correctness, not speed or app behaviour; a release that changes the engine or the app gets a run on a real phone first.

Where numbers come from

Benchmark figures in the docs are measured (12 GB / UFS 4.x test phone, Qwen3-30B-A3B-Q4_K_M and friends). Don't invent or round them silently; if you re-measure, update docs/benchmark-method.md and the README table together. Phrase the test device generically in anything public.

6.1 KiB Raw Blame History