[bong-water-water-bong]

Closes #984.

Root cause

snapshot_state() cut every plane with the per-position count tokens * NKV * HD, including the GDN conv window
and the recurrent state, which are one fixed block per layer that each position updates rather than appends to.
On Bonsai-2-27B (per-position 1024 floats, conv 30,720, recurrent 786,432) that broke the planes on both sides:

• a prompt longer than 30 tokens makes tokens*1024 exceed the conv plane, so restore_state's "must cover the
  positions being restored" test failed and the prefix cache fell back to a full prefill;
• a prompt shorter than 768 tokens made the snapshot keep only tokens*1024 of the 786,432-float recurrent plane
  (23 % at 180 tokens), and the install would then shorten the plane the GDN step reads.

Measured before

On the shipped serve (t2s2-20261103/1bit, :18097) every shared-prefix request answered
[prefix] slot=-1 cached=0 of 849 reason=slot has no usable state (5 of 5 long requests across three windows, on
the ladders recorded in ~/wt/prefill-rows/), and a byte-identical repeat of an 849-token prompt paid its whole
prefill twice: 104.0 s total against an 81.9 s TTFT.

The saved sizes are the fingerprint — one serve, [prefix] saved slot=...:

prompt tokens	slot bytes	393,216*tokens + 7,864,320
180	78,643,200	78,643,200
252	106,954,752	106,954,752
584	237,502,464	237,502,464
704	284,688,384	284,688,384
849	320,471,040	341,704,704

384 KiB/token against the 64 KiB/token of KV the geometry has, i.e. the fixed planes were being grown as if they
were per-position; at 849 tokens the slot is 21,233,664 B below the line, exactly 64 layers * 82,944 floats,
which is those planes reaching their real size.

The change

• npu/forward/include/q4nx_state_planes.h (new, header only, no device): take_per_pos/give_per_pos for the
  KV, take_whole/give_whole for the GDN planes, with the shared-KV empty-layer case kept as it was.
• snapshot_state()/restore_state() use them. A whole plane is accepted only at the size this forward
  allocated, so a truncated snapshot is refused (caller re-prefills) instead of installed short, which would
  leave the GDN step reading past the end of the plane. Net −19/+16 lines in q4nx_forward.cpp.
• tests/q4nx_state_planes_test.cpp (new ctest q4nx_state_planes, header only, runs in CI): the old arithmetic
  reproduced (both symptoms), the round trip at the real plane sizes, the KV per-position behaviour unchanged, and
  the truncated-plane refusal. Wired with the same add_dependencies(npu_sampling_test ...) trick as
  mtp_draft_test, because ci.yml lists every ctest target and this branch cannot edit that file.
• docs/npu.md: a short section stating the two kinds of plane and what the old behaviour cost.

Verification

• tools/copyright.py --check rc 0.

• g++-16 -std=c++17 -Wall -Wextra -Werror tests/q4nx_state_planes_test.cpp green; cc1 mutation checks: carrying
  a whole plane with the per-position cut fails 9 checks, accepting a below-size plane fails 2.

• The changed translation unit compiles with the tree's own NPU flags
  (g++-16 … -DONEBIT_NPU_RUNTIME_OWN -std=c++26 -fopenmp -c npu/forward/src/q4nx_forward.cpp) rc 0.

• Device A/B on this box, one npu-hold ticket, one serve per binary, same 180-token prompt sent twice, greedy
  (evidence in the comment below, serve logs in ~/wt/prefill-rows/):

  binary	request 2 decision	slot bytes	TTFT 1st -> 2nd	text
  shipped t2s2-20261103 (md5 2d66b8ca)	slot=-1 cached=0 of 180 … slot has no usable state	78,643,200	19.26 -> 22.17 s	8c116342… both (full re-prefill)
  planes-only 5b87def	slot=0 cached=179 of 180 … hit	232,783,872	15.63 -> 5.32 s	differs (the second bug)
  this tip	slot=0 cached=179 of 180 … hit	232,652,800	22.90 -> 7.89 s (0.34x)	identical 8c116342…, equal to the shipped binary

  amdxdna errors 0 in both windows, both serves stopped with SIGTERM inside their grace windows.

• Device-free re-check after the rebase onto f63af5f: tools/copyright.py --check rc 0 and
  g++-16 -std=c++17 -Wall -Wextra -Werror -Iapp -Inpu/forward/include tests/q4nx_state_planes_test.cpp prints
  q4nx_state_planes: all checks passed.

No merge, no tag, no deploy: the owner merges.

[github-actions]

PR Reviewer Guide 🔍

(Review updated until commit 8ac86d0)

Here are some key observations to aid the review process:

🎫 Ticket compliance analysis 🔶 984 - Partially compliant Compliant requirements: The prefix cache now correctly restores state by carrying GDN planes whole Root cause addressed: GDN conv window and recurrent state are no longer truncated Hardware measurements confirmed: 5 of 5 long requests no longer re-prefill Tests added in tests/q4nx_state_planes_test.cpp to verify the fix Documentation updated in docs/npu.md to explain the two kinds of planes Non-compliant requirements: None Requires further human verification: None

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

🧪 PR contains tests

🔒 No security concerns identified

⚡ Recommended focus areas for review Incorrect Plane Cutting Logic The snapshot_state function was incorrectly cutting all planes (including GDN conv and recurrent states) with the per-position count, which truncated fixed-size planes. This caused restore_state to refuse snapshots, leading to full prefill on every request. The fix correctly separates per-position planes (kcache_, vcache_) from whole planes (conv_state_, rec_state_) in the snapshot logic. bool Q4nxNpuForward::snapshot_state(int tokens, StateSnapshot& out) const { if (!ready_ || tokens <= 0) return false; const size_t per = size_t(c_.NKV) * size_t(c_.HD); size_t held = 0; // positions the KV caches hold (a layer that shares its KV holds none) for (const auto& layer : kcache_) held = std::max(held, layer.size() / per); if (size_t(tokens) > held) return false; out.tokens = tokens; out.per_pos = per; const size_t want = size_t(tokens) * per; // KV is per position, so a snapshot of the first `tokens` positions keeps exactly that many; the GDN conv // window and recurrent state are fixed-size blocks updated per position, so they go in whole. Cutting them // with `want` truncated them (23 % of the recurrent plane at 180 tokens) and made restore_state refuse the // snapshot -- the prefix cache's `slot has no usable state`, which cost a full re-prefill on every request. q4nx_state::take_per_pos(kcache_, out.k, want); q4nx_state::take_per_pos(vcache_, out.v, want); q4nx_state::take_whole(conv_state_, out.conv); q4nx_state::take_whole(rec_state_, out.rec); return true; } Resume Point Logic The resume point calculation was changed to ensure that snapshots are taken one token short of the prompt, which prevents applying the last token's GDN update twice. This is critical for GDN models where planes are updated in place rather than appended to. // One token short of both the slot and the prompt: the snapshot stops one token short (see the save // below), and the last token is run again for the logits the answer is drawn from. const int want = npu::resume_point(d, ids.size());

[bong-water-water-bong]

READY on the device-free half. What I ran, on 5b87def:

$ g++ -std=c++17 -O1 -I. tests/q4nx_state_planes_test.cpp -o t && ./t
q4nx_state_planes: all checks passed        (rc 0)

Two mutations, both caught — so the test discriminates rather than just passing:

1. give_whole() made to install the shorter prefix (the old put() behaviour): 4 checks fail (a truncated rec plane is refused, a refused restore leaves the plane untouched, a below-size rec plane is refused too, a below-size plane does not resize the destination).
2. give_per_pos() with the coverage test removed: 1 check fails (restore refuses KV that does not cover the positions).

Code read, and the failure direction is the right one: a snapshot written by the old code (conv/rec truncated by tokens * NKV * HD) is now refused by give_whole() and falls back to a full prefill, instead of installing a short plane that the GDN step reads past the end of. take_per_pos() keeps the min(want, size) clamp, and the shared-KV empty-layer case is preserved verbatim in both directions.

Not verified by me: the outcome this is for — [prefix] slot=-1 … slot has no usable state disappearing and the second 849-token prompt not paying a full prefill. That needs the NPU; I take no npu-hold ticket, so the serve A/B in your PR comment is the half I cannot re-derive from here.

[bong-water-water-bong]

Measured on the box (this tip, aa2f5c1), and the run found a second bug worth stating plainly.

The first version of this fix (5b87def, planes only) made the cache hit but answer differently, which is why the
tip now also moves the save point. One npu-hold ticket, one serve per binary, prefix cache on by default, the
same 180-token prompt sent twice, greedy, amdxdna errors 0 in both windows:

binary	request 2 decision	slot bytes	TTFT 1st -> 2nd	text
shipped t2s2-20261103 (md5 2d66b8ca)	slot=-1 cached=0 of 180 reason=slot has no usable state	78,643,200	19.26 -> 22.17 s	8c116342101d013a on both (full re-prefill)
planes-only 5b87def	slot=0 cached=179 of 180 reason=hit	232,783,872	15.63 -> 5.32 s	differs — "\n\nThe capital of France is Paris. The ri" then " h o l e prefix with a recurrent state: "
this tip aa2f5c1	slot=0 cached=179 of 180 reason=hit	232,652,800	22.90 -> 7.89 s (0.34x)	identical 8c116342101d013a, equal to the shipped binary's

So the plane fix alone is not shippable: it turns a silent slowdown into silent wrongness. The second bug is the save
point — the snapshot was taken after the last prompt token and restored at cached - 1 positions, which KV allows
(the per-position plane truncates exactly) and the in-place GDN planes do not (the last token's recurrence update
lands twice). forward-serve now snapshots one token short, splitting the final prefill chunk (one extra launch),
and the resume point comes from npu::resume_point(), so both plane kinds are exact where the prefill resumes and
the last token is re-run for the logits the answer is drawn from — which is what the header's own comment on
restore_state() already described ("a request resumes to just short of its own prompt and re-runs the last token").

The slot grows 78.6 MB -> 232.7 MB because the conv + recurrent planes (154 MB at this geometry) are now carried;
that is the cost of correctness, and it is what the cache_mb budget is for.

resume_point() is unit-tested in the same new ctest (q4nx_state_planes) — three mutations are caught: carrying a
whole plane with the per-position cut (9 failures), accepting a below-size plane (2), resuming at cached instead of
cached - 1 (4). Full suite: CI. No merge, no tag, no deploy from me.

[bong-water-water-bong]

READY at aa2f5c1 (the delta from 5b87def, +78/−19 in app/forward_serve.cpp, app/npu_prefix_cache.h, tests/q4nx_state_planes_test.cpp, docs/npu.md, CMakeLists.txt).

Device-free test, run: compiled tests/q4nx_state_planes_test.cpp + q4nx_state_planes.h + app/npu_prefix_cache.h standalone and ran it → q4nx_state_planes: all checks passed (rc 0). Mutations of resume_point(), both caught:

resume = cached            -> 4 check(s) failed  (all four one-token-short checks)
no min() against the prompt-> 1 check(s) failed  ("a short prompt still resumes one token short of itself")

(A third I tried, dropping the slot < 0 || cached <= 0 guard, is behaviourally identical because the function returns 0 in those cases anyway — not a gap in the test.)

The save point, read for the cases that can go wrong. The split is off + take == n && take > 1 → run take - 1 tokens, then the next iteration has take == 1 and falls through to the save branch at off == n - 1, i.e. the snapshot is taken with the last prompt token still unevaluated, and the same token is then run to produce the logits — exactly the point resume_point() restores to. chunk_max == 1 degenerates to the same thing without splitting, ids.empty() cannot reach the branch, a failed snapshot_state() only skips the insert, and a resumed request (resume > 0) re-saves at n - 1 with the resumed state, which is consistent. insert(ids, …) still keys the slot on the whole prompt while the snapshot covers n - 1, which matches the new cached - 1 convention on the read side.

One operational thing: gh pr view 985 still reports head 5b87def while the branch npu/snapshot-whole-gdn-state is at aa2f5c1 (git ls-remote agrees with the branch). So CI and the re-gate are currently running against the previous commit — worth a push/touch to get the PR's recorded head onto aa2f5c1 before the morning pass reads the checks.

Same limit as before: the serve number (0.34x TTFT on a hit, identical text) is yours — no npu-hold ticket from me.

[github-actions]

Preparing review...

[bong-water-water-bong]

VERDICT: READY at 133278a31 (re-verdict; the head moved and the new commit is a real behaviour fix, not a rebase).

Two commits since my last verdict at aa2f5c1: the branch was rebased onto main f63af5f and 133278a — "make the
prefix cache's snapshot exact, so a hit reuses it instead of re-prefilling" — landed on top. Its content, not the base
move, is what follows.

What I ran at this head:

$ g++ -std=c++17 -O1 -I npu/forward/include -I app tests/q4nx_state_planes_test.cpp -o /tmp/qsp && /tmp/qsp
q4nx_state_planes: all checks passed          # rc 0

The invariant the new commit rests on, which is the thing I went looking for: a slot's snapshot holds cached - 1
positions, and nothing else can insert a slot. Both halves check out on the tree.

• grep -rn "snapshot_state(\|prefix->insert(\|resume_point(" app npu/forward/src finds exactly one save
  (app/forward_serve.cpp:289-291) and one restore (:257), and the save is inside the split-final-chunk branch, so
  there is no second path that could store a snapshot at full length and make resume_point off by one.
• The split is sound by construction: with chunk_max >= 2 the last chunk runs take-1 tokens and off becomes
  n-1; the next iteration has take == 1, so it takes the save branch at off == n-1 and then runs that token — the
  snapshot precedes the token whose logits the answer is drawn from, which is the only exact point for planes that are
  updated in place. With chunk_max == 1 the split never triggers and the same state is saved at off == n-1
  directly. A resumed request (off starting above 0) re-saves at n-1 from a state that was restored to resume
  positions, which is again exact.
• resume_point() clamps on both sides (min(cached, prompt) - 1, 0 when the slot is unusable or the prompt is
  empty), so a slot cannot be restored past the positions it holds and a hit never resumes at the prompt's own last
  token.

Nit: the split's cost is stated as "one launch" and the whole fix is only covered by the header test — forward_serve's
chunk loop has no test of its own, so the "save one short, resume cached-1" pairing is pinned by reading plus the
serve log ([prefix] saved slot=… tokens=…), not by the suite. Worth a serve-level case when this lane next has a
ticket; the failure it prevents (restoring one position too many, which applies a GDN update twice and answers the same
prompt differently — engine #984) is exactly the kind that a log line will not catch.

CI on 133278a31 was still running when I posted this (build/Analyze pending, CodeQL pass, hw_gate skipped).

[bong-water-water-bong]

One cross-check the measured slot size gives for free, and one thing this PR deliberately does not fix.

The planes-only arm's slot was 232,783,872 B at 180 tokens. That is exactly the model, no slack:

plane	bytes
KV, 16 full-attention layers, 180 * 16 * 2 * 1024 * 4	23,592,960
conv, 64 layers * 30,720 * 4	7,864,320
recurrent, 64 layers * 786,432 * 4	201,326,592
sum	232,783,872

So the fix carries each GDN plane at the size this forward allocated (the whole 786,432-float recurrent block per
layer), the per-position part is only the 16 layers that really have a KV cache (128 KiB/token, not the 64 KiB a
reading of the layer table gives), and nothing is double-counted. The old arithmetic is the same model with the
recurrent plane cut to tokens*1024 floats per layer, which is the 393,216*T + 7,864,320 line in the body:
393,216 = 131,072 (KV) + 262,144 (64 layers * 1024 floats of a plane that is not per-position).

Not fixed here, on purpose: conv_state_/rec_state_ are allocated for all NL layers while only the
Q4NX_LAYER_LINEAR ones ever read or write them, so 16 of the 64 layers contribute 50,331,648 + 1,966,080 B of
zeros = 22 % of every slot (and 52 MB of memcpy per save/restore). That is a memory-efficiency follow-up —
allocating those planes per layer type changes the load path, not the snapshot rule, and it is not needed for the
cache to work. Recorded, not bundled.

[github-actions]

Persistent review updated to latest commit 8ac86d0

[bong-water-water-bong]

The 1k row this fix was never gated on: the hit works, and the resumed prefill is bit-identical to a fresh one

The device A/B in this PR was 180 tokens. At 1k the chunking, the two planes' sizes and the resume point all differ, so
it was worth closing on a merged-main build (origin/main f17cc61, build-npu/1bit md5 c65657d0, the logged recipe,
one npu-hold ticket, amdxdna ERROR lines 0):

request	prompt	TTFT	serve's decision
warm	157 tok	fill	no slot shares a prefix -> saved slot=0 tokens=156 bytes=229,638,144
1st	1021 tok	74.38 s	no slot shares a prefix -> saved slot=1 tokens=1020 bytes=342,884,352
2nd	1021 tok, identical	0.434 s	slot=1 cached=1020 of 1021 reason=hit

171x on the second TTFT, and the slot is exactly 1020 x 131,072 B of KV plus 209,190,912 B of conv+recurrent --
the planes carried whole at 1k, matching the byte accounting in the comment above.

Text is not a usable gate here, so the check is the distribution

Both arms emitted completion_tokens: 0 and text_sha256 = e3b0c442... (the hash of ""), twice, on two different
prompt shapes. The logit dump says why rather than leaving it as a mystery: at 1021 tokens the argmax is 248046 (EOS)
at 14.477 against 271 ("\n\n") at 13.729 -- the near-tie already recorded in prefill-bench.py, except here EOS
wins outright, so a text comparison is a comparison of two empty strings.

So the gate compares the raw float32 vocabulary of the first decode position (serve --dump-logits, one process,
the dump copied aside between the two requests), which is TBONSAI-HYBRID.md SS8-H1's criterion:

vocab 248,320   bit_identical true   max_abs_delta 0.0   nonidentical 0   KLD(fresh||hit) 0.0
top5 fresh == top5 hit: 248046 +14.476595, 271 +13.729347, 561 +12.692009, 19241 +12.276074, 4558 +10.509911

All 248,320 logits are bit-identical. That is a stronger statement than the 180-token text identity this PR shipped
with: restoring 1020 positions and replaying token 1020 produces exactly the distribution a full prefill produces, i.e.
the conv window and the recurrent state are installed whole and each in-place update is applied exactly once.

Harness for both runs is ~/wt/prefix-logits-gate.{sh,py} (and a SUFFIX option on ~/wt/prefix-repeat.py, which is
what showed the empty-text trap). No engine change in this comment -- evidence only, on merged main.