Repository navigation
npu: the MTP drafter runnable against the real base, and the head checked against the reference - #811
Conversation
bong-water-water-bong
left a comment
There was a problem hiding this comment.
VERDICT: CHANGES
Two items, both small; the pipeline itself builds and the draft loop is still green.
What I ran (-DONEBIT_NPU=ON -DONEBIT_NPU_RT_DIR=<private runtime>, so the NPU lane is actually compiled — CI does not):
$ ninja mtp_draft_test mtp_batch_all_test mtp_acceptance rc 0
$ ./mtp_draft_test … all checks passed (69 assertions, incl. pos threading:
"commits=12 last_commit=1 accepted=12 drafted=24")
Item 1 — the head-vs-reference number cannot be reproduced from the tree. tools/mtp_acceptance.cpp:162 reads
hidden.bin, emb.bin and out_hidden.bin from --gold DIR, the tree ships none of them, and nothing in the tree writes
them — so --selftest-head, which is the only check of the C++ head's arithmetic, fails with cannot read the golden files under for anyone who is not the author:
$ ./mtp_acceptance --selftest-head --head-exact …/pb16i-exact --head-q4k …/q4_k
cannot read the golden files under
The claim in the body (rel-L1 1.838e-04, max|d| 1.567e-02, cosine 0.999999985) is self-consistent — it lands where the image's
own int16-mantissa limit does (#791 measured 1.56e-04 on the image itself), which is exactly why I want it runnable rather than
taken on trust. Ask: ship the three files (3 × 5,120 floats ≈ 60 KB, and the repo already ships binary goldens under
tests/golden/), or a ~15-line generator built on tools/mtp_head_ref.py head_step so the fixture is derivable from the
reference that defines it.
Item 2 — the branch is behind #807/#809/#810, and the rebase touches the same file as the plan layer. Its own change is
additive: three read-only accessors (lm_head_logits(), embed_row(), last_residual()), one last_residual_.assign(...) in
the step path, and the new npu/mtp/* + tools. What looks like deletion in npu/forward/include/q4nx_forward.h against current
main is #807's plan_/plan_path_/missing_kernel(), which this branch predates. So: rebase needed, and re-verdict after
it, because the accessors will land next to the plan members in one file.
Nothing else is in question: the arithmetic decomposition is the same one #791's artifact and #809's chunk use (signed int4
codes × int16 magnitudes, one exponent per column, the additive min as a host table), forward_base.h's
commit_verify() as snapshot-then-replay is correctly described as "correct, not free" until #798's deferred design is wired,
and the vocabulary projection going through the base model's own lm head is the right call (the drafter has no head of its own —
lm_head_logits() exists for exactly that).
bong-water-water-bong
left a comment
There was a problem hiding this comment.
VERDICT: READY
Correction to my own review: item 1 is closed, and I closed it myself rather than waiting. The head-vs-reference number is
not only reproducible, it reproduces — I generated the missing fixture from the Python reference (the same convention
check_head() already uses) and ran the tool's selftest against it:
# tools/ → the three files mtp_acceptance --selftest-head reads
W, _ = mtp_head_ref.load_q4k_head('…/bonsai2-27b-mtp-head/q4_k')
h, e = np.random.default_rng(0).standard_normal(H), np.random.default_rng(0).standard_normal(H) # unit-RMS, as check_head does
_, out, _ = mtp_head_ref.head_step(h, e, W, pos=3) # pos=3: the position the C++ selftest uses
h.astype('<f4').tofile('hidden.bin'); e.astype('<f4').tofile('emb.bin'); out.astype('<f4').tofile('out_hidden.bin')
$ ./mtp_acceptance --selftest-head --gold <that dir> --head-exact …/pb16i-exact --head-q4k …/q4_k
head selftest vs Python reference: rel-L1 1.933e-04, max|d| 1.758e-03, cosine 0.999999981 rc 0Same order as the body's 1.838e-04 / 0.999999985, well inside the tool's own thresholds (rel-L1 < 5e-3, cosine > 0.9999), and
max|d| differs only because my fixture is a different draw than the author's. So the C++ head really is the reference's
arithmetic at the int16-mantissa limit — the claim stands, checked by me on this box.
That turns my item 1 from an objection into a reproducibility nit, which I am not holding the PR for: the recipe above (or
the three 20 KB files) belongs in the tree, because the next person who wants to check the head should not have to rediscover
the convention. One sentence or one script.
Item 2 stands as a process note, not a defect: the branch predates #807/#809/#810, its own change is additive, and the
apparent deletions in q4nx_forward.h are #807's plan members. When pi-bonsai rebases it I will re-post at the new head (one
line, own-path diff) — the label on this head means "content READY", not "merge without the rebase".
Everything else from my first review holds: NPU lane compiles (rc 0, private runtime, which CI does not do) and mtp_draft_test
is green at 69 assertions with pos threaded through DraftHead::draft.
…cked against the reference
Host-side pieces the device window needs, no device in any of them:
* npu/mtp/mtp_head.{h,cpp} -- the head on the host in the pb16i-exact arithmetic from the container: per
32-K group a signed int4 code and an int16 magnitude, one exponent per column, and the source's additive
min as a host table (y = sum_g (sum_k a_k code) ms 2^eref + sum_g A_g add_g). Same decomposition the
device does, so a draft measured here is a draft the NPU head can reproduce. Its vocabulary projection
goes through the BASE model's lm head (the drafter has none of its own).
* npu/mtp/forward_base.h -- BaseModel over Q4nxNpuForward: step, step_batch_all, last_hidden/last_residual,
and commit_verify() as snapshot-then-replay while the engine has no GDN deferral.
* engine: lm_head_logits(), embed_row(), last_residual() -- the three read-only accessors the drafter needs.
* tools/mtp_acceptance.cpp -- the measurement itself: 8 prompts x 64 tokens, draft k, verify through
step_batch_all, commit the accepted prefix, report acceptance length / hit rate / passes per token, both
hidden-state variants, and a --selftest-head mode that needs no device.
* tools/mtp_head_gemv_harness.py -- freezes the per-matrix fixture (int8 A, the exact int32 C the device must
produce, and the float oracle for the same rows) for the six head matrices that map onto sealed M=16 ELFs.
Checked host-side: the C++ head against the Python reference on a frozen input -- rel-L1 1.838e-04, max|d|
1.567e-02, cosine 0.999999985 (the int16 mantissa, i.e. the same limit as the image itself).
…e window The head's q tensor is 12,288 columns and the sealed QKVFAQ ELF is 6,144 wide, so the harness generator now also emits q_lo and q_hi from the source cell grid (which splits cleanly on the row-block axis) -- two launches, one per ELF, each with its own int32 expectation and float oracle. Every task dir carries its per-column exponent (EREF) so the checker has no special cases. mtp_acceptance gains --prompt-base so a window can run the 8 prompts across several <=3-minute tickets without repeating the same two.
…wn stream The live verify loop cannot run yet: restore_state() refuses a prefix restore on a GDN model (it compares the requested per-position KV length against the conv/rec arrays, which are fixed-size windows rather than per-position state), so snapshot-and-replay is either refused or would advance the recurrence over rejected rows. So the drafter is measured the way the published numbers are: run the base greedily, prime the head's own KV over that stream position by position (which is how the MTP layer is trained), chain k drafts from each position and compare them with the tokens the model actually produced. Same quantity, no base-state surgery. Measured on Bonsai-2-27B, 8 prompts x 64 tokens, k=4: first draft 69.4 %, d2 58.2 %, d3 49.6 %, d4 41.4 %, hit rate 0.694, 3.19 tokens per pass with the post-final-norm hidden. The pre-norm residual variant matches at depth 1 (70.4 %) and collapses after it (2.1 tokens/pass), which settles the plan's open question: the head wants the post-norm hidden. Also here: --residual / --fc-hidden-first (the fusion order is measured, not assumed), --tau 0 for unconditional depth statistics, and step() now records last_residual() as the batch path already did.
6f74773 to
04361f7
Compare
PR Reviewer Guide 🔍(Review updated until commit e32370d)Here are some key observations to aid the review process:
|
restore_state() refuses a prefix restore on a GDN model (it compares the requested per-position KV length
against the conv/rec arrays, which are fixed-size windows), so a speculative verifier had no way back. This
adds the design that needs no way back:
begin_verify(rows, start_pos) snapshots the conv ring + recurrent state ONCE per verify batch and makes
room for each row's recurrence inputs;
gdn_layer_batch() records, per GDN layer and row, the row's conv input and its alpha/beta
gate inputs while the batch runs (the batch itself is unchanged);
commit_verify_rows(accepted) truncates the KV cache to the accepted prefix (position-indexed, so no
state has to be undone), restores the snapshot, and re-applies ONLY the
accepted rows' recurrence through gdn_apply_rows() -- the recurrence alone,
no projections and no GEMMs.
Verified live on Bonsai-2-27B with the MTP drafter: the speculative loop emits **text identical** to plain
greedy decode (k=3/tau=0.6 and k=4/tau=0). Its acceptance is still low in the live loop (0.62 vs the offline
protocol's 2.19) because the drafter's own KV cache is reset every iteration -- the head must be advanced over
the accepted rows, which is the next change; the deferral is what makes that safe to do.
|
Persistent review updated to latest commit cb36628 |
…own step) The MTP layer attends over its own stream, so a drafter that is reset every iteration drafts against nothing -- the offline protocol showed exactly that (3 % acceptance unprimed, 69 % primed). The loop now hands the drafter each accepted row's post-norm hidden (BaseModel::last_hidden_rows, filled by step_batch_all) and the token the base consumed outside the batch, so the head's stream stays in step with the base's; the drafter keeps only its committed prefix and drops the previous chain's provisional positions (KvCache::truncate). Live result on Bonsai-2-27B, text still identical to plain greedy decode: k=4 tau=0 -> accept-len 1.12, 1.00 passes/token; k=4 tau=0.6 -> 0.70, 1.25. That is still well short of the offline protocol's 2.19 (0.69 hit rate), so the remaining gap is in how the live loop feeds the drafter, not in the deferral: next is per-depth accounting in the live path plus priming from the batch's own post-rollback rows.
|
Persistent review updated to latest commit f544032 |
|
Persistent review updated to latest commit e32370d |
bong-water-water-bong
left a comment
There was a problem hiding this comment.
VERDICT: READY
At e32370d9, the rebased head (own paths byte-identical to f5440322 — git diff --stat over npu/mtp/,
npu/forward/{include,src}/q4nx_forward.*, tools/mtp_acceptance.cpp, the harness and the two tests → empty), build green,
state clean.
Everything from my reviews stands, including the one I verified myself: the C++ head reproduces the Python reference at
rel-L1 1.933e-04, max|d| 1.758e-03, cosine 0.999999981 on a fixture I generated from mtp_head_ref.head_step (pos=3) —
i.e. the int16-mantissa limit, not merely "close". mtp_draft_test is green at 69 assertions with pos threaded through
DraftHead::draft, and the new commits were read: commit_verify() now truncates the KV caches to
verify_pos_ + accepted, restores snap_conv_/snap_rec_ and re-applies only the accepted rows through
gdn_apply_rows(0, accepted), which is correct for a recurrence and cheaper than replaying the whole batch.
Two notes to carry, neither blocking:
- ship the head-vs-reference fixture (three 20 KB files, or the ~15-line generator:
default_rng(0)for hidden/emb, then
head_step(..., pos=3)). The check works and its thresholds are tight; it just cannot be run by anyone who is not the
author today. - "deferred GDN commit" is a snapshot-then-replay of the accepted rows, not #798's deferral — the commit subject
overstates it a little. The behaviour is correct; the wording is what a later reader will trip on when they go looking for
the deferred mechanism in the forward (it lives indraft.h's interface, and this adapter satisfies it with a snapshot).
tprs-ready is on. This is the head to merge on.
…row accessor can no longer stop overriding (#1143) * npu/mtp: the forward adapter describes the deferral it uses, and its row accessor can no longer stop overriding Three defects in one small header, all of the same family as the #811 drift: * the file comment said the adapter snapshots the state per verify batch and restores the accepted prefix, because "the forward has no GDN deferral yet". The code has used the deferred path since #805 P6 (begin_verify stores each row's recurrence inputs, commit_verify_rows applies the accepted prefix); the snapshot design is what the deferral replaced, and the comment priced the loop at a state copy it does not take. * the members that comment left behind -- snap_ (unique_ptr<StateSnapshot>) and snap_pos_ -- were dead: one declared and never touched, one assigned and never read. * last_hidden_rows() was the only non-override in the class, and it is the single signature the drafter's advance hangs on: drop the const and it silently stops overriding, reporting no rows at all, which is the failure mode that cost the live loop 1 token/pass for a day. Proof the guard works: with override, removing const from the accessor is now a compile error (`marked 'override', but does not override`), where before it was silent. mtp_acceptance builds clean, ctest -R mtp_draft 1/1, copyright --check clean. * npu: the adapter's comment says what a prefix restore DOES, not that it is refused Review nit on #1143, and correct: `restore_state()` (npu/forward/src/q4nx_forward.cpp:326-341) does not refuse a prefix. It refuses only a plane of the wrong size (`give_whole`) and a `tokens > in.tokens` extension; for `tokens < in.tokens` it restores the KV caches per position to `tokens` and then installs the GDN conv/recurrent planes WHOLE from the snapshot. So a prefix restore is accepted and leaves KV at the prefix with the recurrence still at the snapshot's last position -- exactly the mixed state the deferred commit exists to avoid, and exactly what the sentence was meant to warn about. The comment now says that, with the line reference, instead of claiming a guard that is not there. This PR exists because a comment asserted a design the code had replaced; a comment asserting a safety property that is not implemented is the same defect with the opposite sign, so it is fixed rather than merged with it. Comment only: `ninja mtp_acceptance` rc 0, `ctest -R mtp_draft` 1/1, `copyright --check` clean, no behaviour change (the adapter still never calls snapshot_state()/restore_state()). --------- Co-authored-by: bong-water-water-bong <bong-water-water-bong@1bit.gg>
Host-side pieces the MTP window needs, stacked on #797 (
step_batch_all) and #798 (deferred commit + tau). No devicein any of it, and the head is checked against the reference before any window opens.
npu/mtp/mtp_head.{h,cpp}— the head on the host in the pb16i-exact arithmetic the container holds: per 32-Kgroup a signed int4 code and an int16 magnitude, one exponent per column, and the source's additive min as a host
table (
y = Σ_g (Σ_k a_k·code)·ms·2^eref + Σ_g A_g·add_g). Same decomposition the device does, so a draft measuredhere is a draft the NPU head can reproduce. Its vocabulary projection goes through the base model's lm head — the
drafter has none of its own.
npu/mtp/forward_base.h—BaseModeloverQ4nxNpuForward:step,step_batch_all,last_hidden/last_residual, andcommit_verify()as snapshot-then-replay while the engine has no GDN deferral (correct, notfree; the deferred design is npu: the draft loop commits a verify prefix instead of rolling back, and gates the chain on tau #798's target).
lm_head_logits(),embed_row(),last_residual().tools/mtp_acceptance.cpp— the measurement itself: N prompts × M tokens, draft k, verify throughstep_batch_all, commit the accepted prefix, report acceptance length / hit rate / passes per token, bothhidden-state variants (
--residual), plus--selftest-headwhich needs no device.tools/mtp_head_gemv_harness.py— freezes the per-matrix fixture for the six head matrices that map onto sealedM=16 ELFs (int8 A, the exact int32 C the device must produce, and the float oracle for the same rows), so a device
ticket can report cost and byte-exactness rather than just "it ran".
Checked host-side: the C++ head against the Python reference on a frozen input — rel-L1 1.838e-04, max|d|
1.567e-02, cosine 0.999999985, i.e. the int16 mantissa, the same limit the image itself has.
mtp_draftctest stillgreen (69 assertions) with
posthreaded intoDraftHead::draft(rope cannot guess the position of the first draft).Not merged here.