[bong-water-water-bong]

Refs #1781 (the serve half of "the loop inside forward-serve"; follows #2253, #2255, #2245, #2257).

What

forward-serve --spec-lookup K (off by default; --spec-fallback X adds the fallback rule). A request that is greedy, has no grammar, asks for no logprobs, and runs on a server without --dump-logits decodes through mtp::run_from() after the usual batched prefill. Every other request takes the plain loop, unchanged.

• The drafter, npu/mtp/prompt_lookup.h: PromptLookupHead needs no model. It proposes the up to K tokens that followed the latest earlier occurrence of the stream's last 3, 2 or 1 tokens. That pays where a reply repeats its prompt (code being edited, quoting) and proposes nothing otherwise. It is the first drafter the serve can use without a draft package, and it makes the serve path testable on the host. The DFlash drafter (npu: the DFlash draft as mtp::run()'s drafter (DFlashHead, refs #1781) #2255) plugs into the same call.
• One token path: both loops hand every token to one deliver() (end tokens, grammar, stop sequences, the stop hold-back, on_text). It is the old loop body moved, not rewritten, so speculation cannot treat a token differently.
• The speculative path is used only where it cannot change a token. The verify rows check drafts against the same tokenizer-bounded argmax sample_token() takes when greedy.
• The last token: the reply's last token is not run through the model (Options::consume_last = false), as in the plain loop.
• Verify width: in host-GEMM mode the verify batch is 16 rows. Host GEMMs have no width of their own, and the 1 the forward reports there leaves no row for a draft. On the device it is the forward's own width.
• Logging: one [spec] line per request (rounds, drafts accepted, plain steps, fallback).

forward-serve --spec-draft <dir> (--spec-draft-k K, default the package's block minus one) serves the same path with a DFlash draft package as the drafter (DFlashHead, #2255):

• Load: the package is decoded once and checked against the target before ready (the widths must match, and its target layers must exist).
• Taps: the forward writes those layers into a tap sink on every pass, and the sink is emptied when a request takes the engine.
• Pick rule: the target's head, then DFlash 2's candidate selector when the package has one (dflash2_select(), now one function in npu/dflash_head.h), else the bounded argmax.
• Prefix-cache reuse: a request that resumed from a saved prefix has no tap rows for the positions it skipped. Its head refuses to draft, so it decodes with plain steps: the same text, no speedup.
• --spec-draft and --spec-lookup together are refused. 1bit now links npu/dflash_forward.cpp and npu/dflash_draft.cpp.

Two fixes in npu/mtp/draft.h, found by these tests:

• the one-pass round now calls commit_rows() even when the base keeps no hidden rows. Before, a token-only drafter never saw a commit, so every lookup after the first ran on a stale stream;
• Options::consume_last is new (default true, so existing callers see no change).

Proof

• tests/serve_spec_lookup_test.py (ctest serve_spec_lookup, ONEBIT_NPU builds, where forward-serve exists): a fabricated two-layer qwen3 container with a 256-token byte-level tokenizer, served host-GEMM twice, plain and with --spec-lookup 4. No device is opened: under strace -e openat the whole run shows 0 opens of /dev/accel* or amdxdna. Text, finish_reason and completion_tokens are equal on all 16 requests:

  • 4 prompts × budgets 1, 5 and 24;
  • 3 stop sequences cut out of the plain replies, which end them (finish stop, at 8 to 12 tokens);
  • a seeded temperature 0.8 request.

  The speculative path was taken on exactly the 15 greedy requests, with 6 of 33 drafts accepted on this random model, and never by the plain server or on the sampled request.

  A third server runs --spec-draft on a generation-1 DFlash package the script writes (tests/dflash_draft_test.cpp's synthetic layout, sized to the fabricated target). All 16 replies equal the plain server's. The speculative path was taken on exactly the 15 greedy requests, and every one of two or more tokens drafted (782 drafts in all). The acceptance means nothing on synthetic weights; what the test pins is the equality and the alignment of the tap sink. The three servers together make 0 opens of /dev/accel* or amdxdna under strace -e openat.

• tests/prompt_lookup_test.cpp (ctest prompt_lookup, host):

  • the lookup by hand: the longer suffix wins, the latest occurrence wins, no match gives no proposal, a misaligned position gives no proposal;
  • through mtp::run(), the text is plain decode's, and the head's stream is the committed text;
  • on an answer that repeats, 30 drafts are accepted at k = 7, and the run takes 18 base passes against plain decode's 48.

• mtp_draft [18]: consume_last = false gives the same text and skips exactly the one pass after a reply that ends on the base's own token.

Mutants, each failing a test:

mutant	failing test
no greedy() gate (the sampled request decodes speculatively)	serve_spec_lookup: 3 failures
deliver()'s false ignored (stop sequences no longer end the reply)	serve_spec_lookup: 4
host mode verifies the forward's 1 row (nothing drafted)	serve_spec_lookup: 2
no commit_rows() without hidden rows	prompt_lookup: 9
--spec-draft: tap sink not emptied per request (warm-up's rows misalign every request)	serve_spec_lookup: 2
--spec-draft: tap layers never set	serve_spec_lookup: 2
--spec-draft: a pick that never proposes	serve_spec_lookup: 2

Full ctest of an ONEBIT_NPU build with the device tests excluded, at d9cdfbb: 271 of 271 passed.

Not measured: tok/s on the device. The NPU is off for now. The rows owed are the served 27B against plain decode, on a code-edit prompt and a prose prompt: with --spec-lookup 4 --spec-fallback 2, and with --spec-draft ~/scratch-packs/dflash2-27b.

🤖 Generated with Claude Code

[bong-water-water-bong]

ready for review at eadff3a

[github-actions]

PR Reviewer Guide 🔍

(Review updated until commit d9cdfbb)

Here are some key observations to aid the review process:

🎫 Ticket compliance analysis 🔶 1781 - Partially compliant Compliant requirements: Implement speculative decode on the NPU with a prompt-lookup drafter that proposes up to 7 tokens in one pass The verify stays one pass on the NPU Distribution-preserving sampling is implemented with the prompt-lookup drafter The implementation includes both nemotron3.5-bolt + DFlash draft and qwen3.8-mtp:27b + DFlash 2 The server supports --spec-lookup K and --spec-draft options Non-compliant requirements: The specific parity numbers (2.84 tok/s) are not measured in this PR Requires further human verification: The specific performance numbers (2.84 tok/s) need to be verified on hardware The distribution-preserving sampling behavior needs to be validated with actual model outputs

⏱️ Estimated effort to review: 4 🔵🔵🔵🔵⚪

🧪 PR contains tests

🔒 No security concerns identified

⚡ Recommended focus areas for review Speculative Decode Logic The speculative decode logic in handle_generate_impl introduces a new code path that uses mtp::run_from() with a PromptLookupHead. This path is only taken for greedy requests without grammar or logprobs, but the condition check opts.sample.greedy() && !gram && !want_lp && e.dump_logits.empty() may be too restrictive and could potentially exclude valid cases where speculative decoding should be used. The logic should be verified to ensure it correctly identifies when speculative decoding can be safely applied without changing the output. const bool spec = (e.spec_lookup > 0 || e.draft) && opts.sample.greedy() && !gram && !want_lp && e.dump_logits.empty(); if (spec) { Prompt Lookup Implementation The PromptLookupHead implementation uses a nested loop to find the latest earlier occurrence of a token sequence. This approach may have performance implications for large token sequences or when the token stream is very long, as it performs a O(n^2) search in the worst case. While this is acceptable for a host-only implementation, it could become a bottleneck if the prompt lookup is extended to be used in more performance-critical paths. for (int n = std::min(max_n_, len - 1); n >= min_n_; n--) { const int* suf = stream_.data() + (len - n); for (int i = len - n - 1; i >= 0; i--) { // latest earlier occurrence first bool hit = true; for (int j = 0; j < n && hit; j++) hit = stream_[size_t(i + j)] == suf[j]; if (!hit) continue; for (int j = i + n; j < len && int(out.size()) < k; j++) out.push_back(stream_[size_t(j)]); break; } consume_last Option The consume_last option in Options struct allows skipping the final model pass when the budget ends on the base's own token. This is important for performance optimization, but the implementation in run_from function needs to be carefully reviewed to ensure that it correctly handles the case where consume_last is false and the token count reaches the maximum. The check if (!opt.consume_last && (int)r.tokens.size() >= opt.max_tokens) break; should be verified to ensure it doesn't introduce any edge case bugs. if (!opt.consume_last && (int)r.tokens.size() >= opt.max_tokens) break;

✅ Resolved findings

npu/mtp/prompt_lookup.h:50-58

Prompt Lookup Implementation

The PromptLookupHead implementation in npu/mtp/prompt_lookup.h uses a nested loop to find the latest earlier occurrence of a token sequence. This approach may have performance implications for large token sequences or when the token stream is very long, as it performs a O(n^2) search in the worst case. While this is acceptable for a host-only implementation, it's worth noting that this could become a bottleneck if the prompt lookup is extended to be used in more performance-critical paths.

[bong-water-water-bong]

ready for review at d9cdfbb (adds --spec-draft: a DFlash package as the drafter, same equality test with a third server)

[github-actions]

Persistent review updated to latest commit d9cdfbb

[bong-water-water-bong]

VERDICT: READY at d9cdfbb — the guard is the design: the speculative path is only entered where it cannot change a token, and the test asserts that from both sides.

"Used only where it cannot change a token"

"A request that is greedy, has no grammar, asks for no logprobs, and runs on a server without --dump-logits"

Four exclusions and each is load-bearing:

• greedy — with sampling, drafting is not an optimisation but a different distribution, so it must not engage at all;
• no grammar — a grammar can reject the argmax draft, so the accepted set is not the plain greedy set;
• no logprobs — speculative acceptance changes which steps were run, and logprobs are a per-step observable;
• no --dump-logits — the same argument, for the artifact the whole measurement apparatus reads.

Naming these as preconditions rather than caveats is the difference between a feature that is safe by construction and one that is safe if the user reads the docs.

The test asserts both directions

From the ctest:

• id-equality on the same id prompts across several budgets — "must be equal";
• the spec server must take the speculative path on exactly the greedy requests (its [spec] line), i.e. the arm that would pass trivially if speculation silently never engaged;
• a sampled request must stay on the plain loop — the negative direction, which is the one that catches "we enabled it everywhere and got lucky on the greedy prompts".

That trio is the pattern I have been asking for all night in smaller pieces: equality where it must hold, engagement where it must happen, and abstention where it must not.

And the honest gap is stated, with its reason

"Not measured: tok/s on the device. The NPU is off for now. The rows owed are the served 27B against plain decode…"

Refs #1781, not Fixes, with the owed measurement named and why it is owed (owner's NPU-off order) rather than left as a TODO. That is the correct way to land the serve half while the device is unavailable — and it means nobody later mistakes this for a performance claim.

Two asks

1. state the verify batch width's cost at the boundary — 16 rows in host-GEMM mode is the verify width; when K drafts are fewer than 15, the batch is padded, so the cost of a round with 2 accepts is a 16-row batch. One line on whether the padding is free (host GEMM has no width of its own — which the body already says) makes the economics legible before the device rows land;
2. name the [spec] log line's format in the docs — a supported mode should be observable from a log without reading the source; it also makes the "engagement" assertion reproducible by someone holding only a server log.

[bong-water-water-bong]

VERDICT: READY at d9cdfbb — fresh configure rc 0, prompt_lookup_test builds and runs: 24 PASS, all checks passed, and the arm that carries the claim is visible in the output:

PASS  period 6 k=7: a repeating answer: 30 drafts accepted, 18 passes against plain's 48
INFO  period 0 k=1/3/7: accepted 0 of 3/9/21 drafts, 48 passes against plain's 48

That pair of lines is the whole design argument, and it is a pair rather than a claim: on an answer that does not repeat, the drafter proposes and nothing is accepted, so the run costs exactly plain decode's 48 passes — the feature is a no-op where it cannot help. On a repeating answer, 30 acceptances turn 48 passes into 18. A speculative path that imposed any overhead in the no-help case would be a tax on every request; this one is measurably free.

The gating is the part I would have insisted on, and it is stated as a closed list: a request decodes through mtp::run_from() only if it is greedy, has no grammar, asks for no logprobs, and the server has no --dump-logits. Every other request takes the plain loop unchanged — which is the right direction for a first serve integration, because each of those four conditions changes what a verify pass must do (grammar and logprobs change the acceptance criteria, dump-logits changes what the batch must leave behind).

Three things I want on the record as right:

1. PromptLookupHead needs no model — the first drafter the serve can use without a draft pack, proposing tokens that followed the latest earlier occurrence of the last 3/2/1 tokens. That is exactly the drafter that can never be wrong about the target's distribution, because the target still verifies every draft.
2. The serve test asserts the speculative path was taken on exactly the 15 greedy requests out of 4 prompts × 3 budgets plus stops and a seeded temperature request — a count, not "it worked". A feature that quietly stopped engaging would fail that line.
3. Mutants table — each mutant fails a named test, which is the property that makes 24 arms mean something rather than being a suite that cannot fail.

Two asks:

1. consume_last = false and the nine bare INFO lines are the two places a reader cannot tell "nothing to do" from "something got disabled". The INFO lines are fine as information, but one arm asserting the reason for zero acceptances on the non-repeating prompt (the drafter proposed and the target rejected, versus the drafter proposed nothing) would stop a future gating bug from looking like a normal no-op. The 3/9/21 draft counts already distinguish them; an assertion saying so would make it a check.
2. State the cost on a real pack, even roughly: the feature is free when it cannot help on this synthetic two-layer model; the number a lane will want before enabling it is what the lookup plus verify costs on a real 27B with K=7 when 0 drafts are accepted. If that is unmeasured, say so in the PR rather than letting "off by default" imply "and cheap when on".

Scope of my check: fresh configure rc 0, built and ran prompt_lookup_test (24 PASS, including the 30-accepted/18-vs-48-passes arm). Did not run serve_spec_lookup_test — it needs an ONEBIT_NPU build and the private runtime, which this box cannot configure.

[bong-water-water-bong]

VERDICT: READY d9cdfbb

The serve integration is gated and tested the way an optimisation of the decode loop has to be, and the two properties that matter are both asserted rather than assumed.

Output equivalence, token for token. tests/serve_spec_lookup_test.py's own contract line is "forward-serve --spec-lookup returns the plain loop's reply, token for token", and it drives three servers — plain, --spec-lookup 4, and --spec-draft on a package it writes itself — over the same id prompts, several budgets and stop sequences cut out of the plain reply. That last part is the one I would have asked for: stop sequences are where a speculative path most easily diverges, because the tokens after a stop are computed by the draft rather than the target, so a comparison that did not exercise them would prove little.

Path selection, not just equality. The test also requires the spec server to take the speculative path on exactly the greedy requests and to draft. That is the vacuity guard, and it is the same lesson as #2247 an hour ago: two servers agreeing proves nothing if the second never speculated. Here a silent non-take fails the test instead of passing it.

The gate itself is a list of correctness reasons rather than conveniences: greedy (acceptance is exact only against the target's argmax), no grammar (its mask is per token and per state), no logprobs (the spec path has no per-token distribution to report), no --dump-logits. Each condition excluded is one where speculation would change the answer rather than the speed, and --spec-lookup is off by default, so no existing request changes behaviour on merge.

And the drafter choice makes this reviewable during a freeze: PromptLookupHead needs no model at all — it proposes what followed the last 3/2/1 tokens earlier in the stream — so the whole loop, its fallback and its accounting run on the host with a fabricated tokenizer and a synthetic DFlash package. That is why I can read this carefully today instead of deferring it.

One nit, not blocking: the test builds a "generation-1 DFlash package ... after tests/dflash_draft_test.cpp's synthetic one", i.e. the format now has two hand-built producers. If a third appears, a shared helper would stop them drifting apart in a way that only shows up when one of them is wrong.

[bong-water-water-bong]

VERDICT at d9cdfbb: READY

The safety claim is structural, and I verified it at the two call sites rather than taking the description's word: line 642 in the speculative loop and line 670 in the plain loop both call the same deliver(t), so end tokens, grammar and stop sequences see every token identically whichever path produced it. That is the property that makes speculation safe here — not a statistic about acceptance rates.

Why this is the right kind of speculative decoding to build first: prompt/ngram lookup drafts tokens that literally occurred earlier in the prompt, and every draft is then verified against the model and only the matching prefix accepted. So there is no quality trade-off at all — acceptance is exact by construction, and a wrong draft costs one verify batch and no correctness. A learned drafter has to earn that property; this one has it for free.

Three things I checked specifically:

• The test exercises the case that would break — the body's "stop sequences cut out of" is real: the same prompts go to both configurations with stop sequences in play, not just a happy-path token count.
• [spec] names which drafter ran (dflash vs lookup) alongside rounds/accepted/plain/fallback. A perf line that said only "spec on" would make the two paths indistinguishable in a log.
• Off by default (--spec-lookup 0), which is right for anything that changes the token loop: it must be enabled deliberately to change behaviour.

One thing worth a follow-up test rather than a blocker: the stop-sequence interaction is covered, but a grammar-constrained request run with and without --spec-lookup compared byte-for-byte is the strongest possible statement of "cannot change a token", and there is no such row yet. The shared deliver() makes it very likely to pass; a row would make it known rather than likely.

[bong-water-water-bong]

VERDICT: READY at d9cdfbb with one ask — three of the four guards around a correctness invariant have no arm, and the invariant is the whole reason this path is allowed to exist.

The claim that carries the PR

// engine#1781: the speculative path, only where it cannot change a token.
(e.spec_lookup > 0 || e.draft) && opts.sample.greedy() && !gram && !want_lp && e.dump_logits.empty()

Four conjuncts, and the code comments the reasoning rather than just the mechanism: greedy means each token is the bounded argmax, so a draft that the verifier rejects cannot change what the plain path would have produced. That is the correct argument, and the fix that makes it true elsewhere is the right shape too — "one token path: both loops hand every token to one deliver()", so end tokens, grammar, stop sequences and the stop hold-back cannot diverge between the two loops. A speculative path that had its own delivery code would eventually differ from the plain one in a case nobody tested.

What is measured, and it is the part I trust most

4 prompts x budgets 1, 5, 24 (15 greedy requests) + 1 seeded temperature 0.8 request
  speculative path taken on EXACTLY the 15 greedy requests — never on the sampled one
  texts equal to the plain server on every request
  6 of 33 drafts accepted on this random model        <- honest: a random model accepts little
  3 stop sequences cut out of the plain replies        (finish "stop", at 8 to 12 tokens)
prompt_lookup_test:  30 drafts accepted at k = 7 on a repeating answer; 18 base passes vs plain decode's 48
mtp_draft [18]:      consume_last = false gives the same text and skips exactly the one pass after a reply ending on a stop

"exactly the 15" is the right form: it is a statement that can fail in both directions — too many requests and the guard leaked, too few and the path never ran. And publishing "6 of 33 accepted" on a random model, rather than a favourable prompt, is the number that tells a reader what this costs when the drafter is guessing.

My ask: !gram, !want_lp and e.dump_logits.empty() are unexercised

The test varies greedy vs sampled, budgets and stop sequences. It never sends a request with a grammar, with logprobs, or with dump_logits set. So three of the four conjuncts could be deleted and every arm above would still pass — and the consequence of deleting any of them is a changed token on a request whose contract says it cannot change, which is precisely the failure this guard exists to prevent.

The arms are small, because the assertion is the same shape as the sampled one:

• a request with a grammar → spec path not taken, text equals plain (a grammar constrains the argmax, so the verifier's accept is no longer "canonical");
• a request asking for logprobs → not taken (the verifier would have to report the draft's logprobs, which it does not);
• dump_logits set → not taken (the dumped rows would be the verify batch's, not the plain path's).

The pattern already exists in the file — the sampled request asserts "[spec]" not in logs for that server — so this is three more entries in the same table, not new machinery. I would rather see the guard's four conjuncts each kill the path once than see another budget variant.

Small notes

• --spec-draft and --spec-lookup refused together — right call: two drafters on one request is a case with no measured answer, and refusing names that.
• consume_last default true — "so existing callers see no change" — is the correct default for a new option, and the mtp_draft [18] arm pins the one pass it saves.
• Prefix-cache reuse ("a request that resumed from a saved prefix has no tap rows for the positions it skipped") is the subtlest line in the description and it is handled explicitly rather than discovered later. Worth a sentence in the reply about what the fallback head does at those positions, because that is the case a future reader will break.