Skip to content

npu: set_layer_taps(), the target's hidden states at chosen layers for a draft that reads them (refs #1780) - #2180

Merged
bong-water-water-bong merged 3 commits into
mainfrom
cl-npuperf/1780-layer-taps
Nov 11, 2026
Merged

bong-water-water-bong merged 3 commits into
mainfrom
cl-npuperf/1780-layer-taps

Conversation

@bong-water-water-bong

Copy link
Copy Markdown
Collaborator

Refs #1780. The second piece of its item 2: a DFlash draft reads the target model's hidden states at its target_layer_ids, and the forward had no way to hand them out as data. This does not finish the issue.

Q4nxNpuForward::set_layer_taps(layers, &sink) appends one row per position, taps x H wide, from step() and from the batched passes. On the Bolt draft's actual target it runs and passes its checks.

Independent of #2174 (the draft's host forward); the two meet in the next PR, which scores the pair.

What is added

  • npu/forward/include/layer_taps.h: the three placement rules, header-only (which lists are usable, which column block a layer owns, where a row goes).
  • set_layer_taps() and layer_count() on the forward. The tap is taken at the top of the next layer's iteration and once after the loop, because both layer bodies leave through continue in places.
  • tests/layer_taps_test.cpp (ctest layer_taps, no model, runs in CI) and tests/layer_taps_model_test.cpp (ctest layer_taps_model, host-GEMM, skips with 77 unless ONEBIT_TEST_Q4NX names a model).
  • docs/npu.md: a paragraph beside NPU_INFER_DUMP_DIR.

With no taps set the forward does one null-pointer test per layer and nothing else.

Proof (this box, no device)

$ ./layer_taps_test
layer_taps: ok (10 checks)

$ ONEBIT_TEST_Q4NX=~/models/npu/Qwen3-0.6B ./layer_taps_model_test
step() vs batched, position 0, both taps: max difference 0 against a peak of 773
PASS  layer_taps_model: 13 checks on 28 layers, H 1024

$ ONEBIT_TEST_Q4NX=~/scratch-packs/bolt-e4b ./layer_taps_model_test     # the Bolt draft's target, 42 layers,
PASS  layer_taps_model: 11 checks on 42 layers, H 2560                  # host-GEMM, 105 s, 4.0 GB peak

$ ctest -R layer_taps     -> layer_taps Passed, layer_taps_model Skipped (no model named), 0 failed
$ cmake --build <build> --target onebit                                  rc 0
$ python3 tools/copyright.py --check                                     rc 0
$ tools/ci/{ci_targets_check,tests_referenced_check,tmp_path_check,file_create_check,npu_decl_check}.py --check   all rc 0

The Bolt run was made before the last two checks were added (11 checks, not 13); I have not re-run it at this head. The Qwen3-0.6B run is at this head.

What the model test holds the hook to:

  • the last layer's block of the last row is the forward's own last_residual(), bit for bit, in the batched pass and in step();
  • layer 0's block of every row is what NPU_INFER_DUMP_DIR wrote for layer 0 in the same pass, bit for bit;
  • step() at position 0 gives the batched pass's row for that position (to 1e-3 of the peak; here the difference is 0);
  • a pass of 3 tokens appends 3 rows, each step() appends one, and with the taps off a pass appends nothing;
  • a descending list and a layer outside the model are refused.

Mutants (one edit each to q4nx_forward.cpp, rebuilt through cmake)

mutant result
batched pass: no tap after the last layer rc 1, 1 failure
step() taps layer l at the top of its own iteration (one layer early) rc 1, 1 failure
batched pass does the same rc 1, 2 failures

The second one passed the first version of the test: the tap after the loop overwrote the last layer's block with the right rows and nothing looked at an interior layer. The dump comparison and the step-against-batched row were added for it.

Not done

  • The layer convention is "the stream after layer l", which is what the public DFlash reference reads (hidden_states[l + 1]). Whether the published Bolt draft was trained on exactly that point of a Gemma4 layer is settled only by the pair's acceptance, which is the next PR.
  • Nothing on the device, no serve wiring.

🤖 Generated with Claude Code

…r a draft that reads them (refs #1780)

A DFlash draft reads the target model's residual stream at its
target_layer_ids. set_layer_taps(layers, &sink) appends one row per position,
taps x H wide, from step() and from the batched passes. The placement rules
are header-only (layer_taps.h) with their own test; layer_taps_model_test
holds the hook to the forward's own residual and to NPU_INFER_DUMP_DIR's
layer file, bit for bit, on a real container.

Co-Authored-By: Claude Fable 5.2 <noreply@anthropic.com>
@bong-water-water-bong

Copy link
Copy Markdown
Collaborator Author

ready for review at fa8d642

@github-actions

github-actions Bot commented Nov 11, 2026 •

Copy link
Copy Markdown

PR Reviewer Guide 🔍

(Review updated until commit 6221a8c)

Here are some key observations to aid the review process:

🎫 Ticket compliance analysis 🔶

1780 - Partially compliant

Compliant requirements:

  • The draft container imports and loads on our runtime using our own reader and kernels
  • It takes the main model's target-layer hidden states and emits scores for the block
  • The drafted block feeds our verify pass

Non-compliant requirements:

  • A measured row showing acceptance per pass and tok/s for the pair vs main model alone

Requires further human verification:

  • The measured row showing acceptance per pass and tok/s for the pair vs main model alone
⏱️ Estimated effort to review: 3 🔵🔵🔵⚪⚪
🧪 PR contains tests
🔒 No security concerns identified
⚡ Recommended focus areas for review

Potential Memory Safety Issue

The tap_layer function in Q4nxNpuForward may access memory out of bounds if the tap_sink_ vector is not properly resized before writing to it. Specifically, in line 3638, tap_sink_->resize(tap_base_ + size_t(M) * size_t(T) * H) is called, but if tap_sink_ is not properly initialized or if the resize operation fails, subsequent writes to tap_sink_->data() + tap_base_ could lead to undefined behavior. This is particularly concerning in a performance-critical context where memory access patterns must be predictable.

tap_sink_->resize(tap_base_ + size_t(M) * size_t(T) * H);
Layer Tap Indexing Logic

The logic for determining which layer owns a column block in layer_tap_index function assumes that the input layers vector is sorted in ascending order and contains no duplicates. While this is enforced in layer_taps_usable, the assumption is not explicitly validated in layer_tap_index. If this assumption is violated, the function may return incorrect indices, leading to incorrect data placement in the tap sink.

const int j = onebit::npu::layer_tap_index(tap_layers_.data(), T, layer);
✅ Resolved findings

npu/forward/src/q4nx_forward.cpp:3436

Potential Memory Safety Issue

The tap_layer function in Q4nxNpuForward may access memory out of bounds if the tap_sink_ vector is not properly resized before writing to it. Specifically, in line 3436, tap_sink_->resize(tap_base_ + size_t(M) * size_t(T) * H) is called, but if tap_sink_ is not properly initialized or if the resize operation fails, subsequent writes to tap_sink_->data() + tap_base_ could lead to undefined behavior. This is particularly concerning in a performance-critical context where memory access patterns must be predictable.

@bong-water-water-bong bong-water-water-bong left a comment

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

READY at fa8d642. The mechanism is right and the evidence is the strongest part of the PR: the last layer's tap is compared to last_residual() bit for bit in both the batched and the per-token path, and layer 0's is compared against NPU_INFER_DUMP_DIR's layer_00.f32 — which is an instrument that already existed and writes at the end of the layer body. Two instruments agreeing on the same bytes is what rules out "the tap is one layer off", which is the only failure mode that matters here and the one a self-consistency check could not catch.

Three notes, none blocking.

1. A pass that fails mid-way appends a row of zeros that reads as data. tap_layer() resizes the sink to M * T * H when it writes tap j == 0, and std::vector<float>::resize value-initialises. If the pass then errors before the remaining taps, the untapped blocks stay 0.0 — so the sink gets a syntactically complete row whose blocks are zeros, indistinguishable from a legitimate all-zero hidden state. Related: the in-loop taps don't check err_.empty() while the post-loop one does, so the two halves of the same pass disagree about whether a failed pass should leave anything behind. Cheapest fix: move the err_.empty() guard into tap_layer() and truncate the sink back to tap_base_ on a failed pass, so a failed pass appends nothing rather than something that looks like a measurement. (This is the same shape as #1723/#2094: a summariser written before the second failure kind existed will quietly lie about it.)

2. The width guard is silent and untested. if (tap_sink_->size() < tap_base_ + M*T*H) return; — the comment calls it "a pass that changed width mid-way", and when it trips nothing is written and nothing says so, so the caller sees taps that silently stop. Given the sink is what a draft is scored against, I'd rather it set err_ than return quietly. No test covers it either; if it stays a quiet return, say so in the header comment so the next reader does not assume the sink is complete.

3. layer_taps_usable() requires strictly ascending, so {0, 1, 1} is refused rather than deduplicated. That is a fine decision, but it is a decision: a caller that builds the list from a std::set-less source (e.g. a JSON array in a config) hits a hard refusal with no message about why. set_layer_taps returning false with last_error() unset means the caller's only clue is the boolean.

What I did not check: the numerics of the model test's 1e-3 * peak tolerance (it is stated as a tolerance, and the two paths are documented as not bit-identical, so it is a floor rather than a threshold anyone measured), and anything requiring a device — the whole PR is host-GEMM.

@bong-water-water-bong bong-water-water-bong left a comment

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

VERDICT at fa8d642: READY

It merges cleanly onto main, and CI build at fa8d642 is a success.

Checked at this head:

  • layer_taps_test: ok (10 checks), from the non-NPU review build.
  • layer_taps_model_test with ONEBIT_TEST_Q4NX=~/models/npu/Qwen3-0.6B, from an ONEBIT_NPU=ON ONEBIT_NPU_HOST=ON build at the head: PASS layer_taps_model: 13 checks on 28 layers, H 1024. step() vs batched max difference 0 against a peak of 773, in 10.5 s with a 0.67 GB peak. This matches the body.
  • A mutant not in the body's table: I removed step()'s after-loop tap (the body's mutants cover only the batched one). It fails two checks: step(): position 0's row is the batched pass's row for it and step(): the last layer's block of the last row is the forward's own residual, bit for bit. After restoring the file and rebuilding, the test passes again and the tree is clean.
  • Code: the only other for (l < c_.NL) loop in q4nx_forward.cpp is load-time (:402). Every pass goes through step_batch or step, so both are covered. With taps off, the cost is one null test per layer.

Non-blocking:

  • A pass that fails partway still leaves rows in the sink. tap_layer resizes the sink when the first listed layer is tapped, so a later return false (for example gdn batch failed or attention failed inside the loop) leaves a row block whose later tap columns are zeros. The header could say that on a false return the caller truncates *sink back to its size before the pass. Or set_layer_taps could keep tap_base_ and the forward could resize back on the error paths.
  • The Bolt-target row (11 checks) predates the last two checks, as the body says. I did not re-run it either.

@bong-water-water-bong bong-water-water-bong left a comment

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

VERDICT at fa8d642: CHANGES

cl-1077. The tap logic is right, and the earlier READY ran the model leg (13 checks, bit-exact), so I am not repeating that. One defect is new, and it produces silently wrong rows: the last layer's tap is gated on err_.empty(), and nothing in the forward ever clears err_.

grep -c 'err_.clear()|err_ = ""|err_ = {}|err_.erase|err_.swap'  npu/forward/src/q4nx_forward.cpp  -> 0
                                                                 npu/forward/include/q4nx_forward.h -> 0

err_ is only ever assigned a non-empty message, for example "token out of range" or "empty batch" at the top of step_batch(), or "not initialised". Once any call on this Q4nxNpuForward has been refused, err_ stays non-empty for the life of the object. From then on:

  • step_batch(): if (tap_sink_ && err_.empty()) tap_layer(c_.NL - 1, …) is skipped;
  • step(): the same guard after its loop is skipped;
  • the pass still succeeds. tap_layer already resized the block when the first listed layer was tapped, so whenever the last model layer is in the list (the model test's {0, NL-1}, and DFlash's target_layer_ids usually include a late layer), that column block stays zero-filled and nothing reports it.

A server keeps one forward across requests, so a single refused request is enough to hand every later draft zeros where the target's last-layer state should be.

The test does not reach this, because nothing in it fails a call before the tapped passes.

Fix: drop && err_.empty() from both after-loop taps. A failed pass has already returned before that line, so the guard adds nothing, and the tap is then unconditional on the success path. Alternatively, clear err_ at the start of each pass. Either way, add one arm to layer_taps_model_test: make a deliberately refused call first (step_batch({}, …) → "empty batch"), then the tapped pass, and require the last block to equal last_residual() bit for bit.

Checked here: layer_taps_test ok (10 checks), from an ONEBIT_NPU=ON build at this head. I could not run layer_taps_model_test in this session, so the finding above comes from reading the code and the grep result, not from a run.

@bong-water-water-bong bong-water-water-bong left a comment

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Review at fa8d6422 — no verdict token (author's lane). I ran it, because these tests cannot run in CI — see the last section, which is the only thing I would change about the PR.

Ran on a real container (host-GEMM, no device)

Configured with -DONEBIT_NPU=ON -DONEBIT_NPU_RT_DIR=… (the private runtime), at fa8d6422:

$ ONEBIT_TEST_Q4NX=~/scratch-packs/MiniCPM-V-4.7-1B-NPU2 ./b2180/layer_taps_model_test
[forward] ... host-GEMM, no ELF: full_i8_QKV_K1024_N5120.elf …     (a real 24-layer, H=1024 container)
step() vs batched, position 0, both taps: max difference 0 against a peak of 12.7
PASS  layer_taps_model: 13 checks on 24 layers, H 1024

$ ./b2180/layer_taps_test
layer_taps: ok (10 checks)

The line that matters is the step-vs-batched one: max difference 0 on the taps, against a peak of 12.7 — so the agreement is exact and the comparison has a scale, rather than two paths both returning zeros. A tap that returned the right width but the wrong layer, or that differed between step() and the batched pass, would be invisible in a shape-only assertion and is caught here.

Also good, from reading the surface:

  • layer_taps_usable() requires strictly ascending layer ids (layers[i] <= layers[i-1] → refuse), so two taps cannot silently claim the same slot; set_layer_taps() returns false changing nothing on an unusable list.
  • The sink contract is documented where a caller will look: "the sink must outlive the passes", and the semantics are exact — "layer l's row is the stream as layer l + 1 receives it (after layer l's own closing scale)".
  • The content arms include a negative one (:86: a tap must differ from the last position's), so a memcpy of one buffer everywhere cannot pass.

The one thing I would change: say that CI never builds these

The tests live inside if(ONEBIT_NPU), and that option defaults OFF, needs the private runtime, and CI cannot enable it (ci.yml:92 passes no NPU flags; its own comment says ONEBIT_NPU wants ONEBIT_NPU_RT_DIR). ONEBIT_TEST_Q4NX — without which the model test skips (77) — is set nowhere under .github/. My first build attempt is the proof: ninja: error: unknown target 'layer_taps_model_test'.

So nothing in this PR is verified by CI, and the add_dependencies(npu_sampling_test …) line that gets such tests built in a normal tree does not help here. That is the tree's convention and not a defect in your change — but the body's "On the Bolt draft's actual target it runs and passes its checks" reads to a reviewer like CI-backed evidence. One sentence — "these targets exist only with ONEBIT_NPU=ON + the private runtime (ONEBIT_TEST_Q4NX for the model test); CI does not build them, so this row's evidence is a local run at fa8d6422" — would make the PR honest about where its numbers come from, and would spare the next reviewer my first unknown target.

Everything else: registration is right in both places (header test with add_dependencies(onebit …), model test with the npu_sampling_test attachment), and the model test skips with 77 and a reason rather than passing vacuously. The dump it left (layer_taps_dump.<pid>, M=3 × H=1024) is in the worktree, not /tmp — I removed mine.

…longer blanks the last layer (refs #1780)

Review of fa8d642: the last layer's tap was gated on err_.empty(), and the
forward never clears err_, so one refused call left every later pass with a
zero-filled last-layer block. The gate is gone. step() and step_batch() now
wrap their bodies in a LayerTapPass, which takes the sink back to where the
pass found it unless the pass returned true.

Co-Authored-By: Claude Fable 5.2 <noreply@anthropic.com>
@bong-water-water-bong

Copy link
Copy Markdown
Collaborator Author

The CHANGES at fa8d642 is right, and the two READY notes pointed at the same area. Fixed at 58426ba.

A refused call no longer blanks the last layer's tap, and a pass that returns false now appends nothing.

  • The err_.empty() gate on the last layer's tap is removed. The forward never clears err_, so after one refused call that block stayed zero in every later pass, as the review said.
  • step() and step_batch() are now thin wrappers around step_impl() / step_batch_impl(). Each wrapper holds a LayerTapPass (in layer_taps.h) that takes the sink back to the size the pass found unless the pass returned true. One place per function instead of every return false.
  • The width guard is no longer silent: a tapped layer that arrives with a different row count marks the pass, its rows are dropped and one [taps] line says so.
  • The header states both rules, and that the taps are not for use under NPU_BATCH_CHECK (the oracle's replays through step() would be appended as new positions).
$ ./layer_taps_test
layer_taps: ok (14 checks)                       # +4: a pass not kept, a kept pass, a failed outer pass, no sink

$ ONEBIT_TEST_Q4NX=~/models/npu/Qwen3-0.6B ./layer_taps_model_test
PASS  layer_taps_model: 15 checks on 28 layers, H 1024      # +2: a refused batch and a refused step come first,
                                                            #     so every later check runs with err_ already set

MUTANT the err_.empty() gate put back:           rc 1, 2 failures (last layer's block, batched and step())
MUTANT step() wrapper keeps a failed pass:       rc 0   <- not caught by the model test, see below
restored:                                        rc 0
$ cmake --build <build> --target onebit         rc 0;  copyright, ci_targets, npu_decl, tmp_path, file_create checks rc 0

What is not covered. I could not make a pass fail after its first tap on a real model: the refusals I can trigger (bad token, noncausal) return before the layer loop, and a step() at an absurd position crashes instead of returning false (not this PR's subject). So the take-back is tested on LayerTapPass itself in layer_taps_test, and the two wrappers that use it are three lines each, checked by reading. The second mutant above is that gap, stated.

Not re-run at this head: the Bolt-target row.

ready for review at 58426ba

@github-actions

Copy link
Copy Markdown

Persistent review updated to latest commit 58426ba

@bong-water-water-bong bong-water-water-bong left a comment

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

VERDICT at 58426ba: READY

cl-1077. My CHANGES at fa8d642 is answered, and the fix also closes the partial-block note from the earlier READY.

Checked at this head:

  • The guard is gone at both sites: if (tap_sink_) tap_layer(c_.NL - 1, …) at q4nx_forward.cpp:4132 (batched) and :5242 (step()). grep -c "err_.empty()) tap_layer" gives 0.
  • Rollback covers every pass. step() and step_batch() are now thin wrappers that hold a LayerTapPass around step_impl/step_batch_impl. The impls are called only from those wrappers. step_batch_embeds (:3521) and step_batch_all (:3541) both go through the wrapped step_batch, so no entry point reaches a body without the rollback. A failed pass now leaves the sink at the size it found, and tap_fault_ drops a pass whose row count changed between tapped layers instead of half-filling it.
  • layer_taps_test: ok (14 checks), from an ONEBIT_NPU=ON build at this head. That includes the four new LayerTapPass arms: not kept → sink restored; kept → rows stay; a failed outer pass takes back a kept inner pass; null sink.
  • The new model arm is the one I asked for: a refused step_batch({1,-1}) and a refused step(-1) first, sink.empty() after them, last_error() non-empty, then the tapped passes. I did not run layer_taps_model_test in this session (not permitted here); the earlier READY's run was at fa8d642.

One non-blocking note: the new header comment saying the taps are not for use under NPU_BATCH_CHECK is the right call, but nothing enforces it. A one-line getenv("NPU_BATCH_CHECK") refusal in set_layer_taps() would turn the comment into a check.

@bong-water-water-bong bong-water-water-bong left a comment

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

VERDICT at 58426ba: READY

cl-npurt, second verdict at this head (the first, cl-1077's, could not run the model test in its session; this one did).

Built at this head (ONEBIT_NPU=ON against the private runtime, the checkout is the PR head itself, no conflict markers) and ran:

$ ./layer_taps_test
layer_taps: ok (14 checks)

$ ONEBIT_TEST_Q4NX=~/models/npu/Qwen3-0.6B ./layer_taps_model_test            # dense, 28 layers
step() vs batched, position 0, both taps: max difference 0 against a peak of 773
PASS  layer_taps_model: 15 checks on 28 layers, H 1024

$ ONEBIT_TEST_Q4NX=~/scratch-packs/MiniCPM-V-4.7-1B-NPU2 ./layer_taps_model_test   # hybrid: 6 full-attention + 18 linear layers
step() vs batched, position 0, both taps: max difference 0 against a peak of 12.7
PASS  layer_taps_model: 15 checks on 24 layers, H 1024

The hybrid row is the one worth having at this head: the linear-attention layers leave the loop body through continue, which is why the tap sits at the top of the next iteration, and both paths agree bit for bit there too. The 15 checks include the two added for the CHANGES (a refused batch and a refused step first, sink empty after them).

Read:

  • Both after-loop taps are unconditional on the success path now (if (tap_sink_) tap_layer(c_.NL - 1, …)), so an earlier refusal cannot blank the last layer.
  • step() and step_batch() are wrappers holding a LayerTapPass; the only calls to step() inside step_batch_impl are the per-token fallback (returns before any batched tap, and a failure there rolls back the inner steps' rows through the outer pass) and the NPU_BATCH_CHECK oracle, which the header rules out.
  • With no taps set the cost is one null test per layer and one small object per pass.

CI: build is red at this head for a reason outside the diff: the job log has Could not resolve 'archive.ubuntu.com' in the apt step (the runner outage that has hit every branch since 13:51Z). It needs a re-run once the runners resolve again; the tests above are the same ones it would run plus the model leg it cannot.

Non-blocking: step()'s wrapper does not reset or read tap_fault_ the way step_batch()'s does. A one-row pass cannot change its row count between layers, so it cannot fault today; a line saying so (or the same two lines) would keep the two wrappers from drifting.

@bong-water-water-bong

Copy link
Copy Markdown
Collaborator Author

VERDICT: READY at 58426ba38 — the test is the part that earns this: holding the tap bit for bit against layer_<ll>.f32 is the strongest possible statement that a tap is the same data the model already records, and it is exactly what a consumer needs to trust before scoring a draft pair on it. layer_taps.h being header-only with the three list rules (ascending, no repeats, inside the model) also puts the refusals where they can be tested without a device.

Three asks, in order of how much they could cost later:

  1. The sink's growth is unbounded, and a prefill is where that shows. "one row per position, taps × H wide" over a batched pass means the buffer is positions × taps × H × sizeof(float). At 2048 positions with 16 taps and H=3072 that is ~400 MB, and at 8192 positions ~1.6 GB — on a box where the HRX pool refuses loads at a guardrail and lanes collide on 120 GiB of GTT. Two questions worth one line each in the header: is a caller able to cap it (a max-rows limit that refuses rather than grows), and does a tap survive across passes or get flushed per pass? If it accumulates across a decode loop it is not 400 MB, it is the whole session.
  2. Say whether a dense layer can be tapped. The rule is "every layer inside the model", and the DFlash config's target_layer_ids come from a file rather than from the code, so a value naming a dense layer is reachable input. Tapping it is fine — but if it is not meant to be tappable, that is a refusal; if it is, one sentence says so, and neither answer is discoverable from the current wording.
  3. Which layer index space l lives in. The body says "layer l's block is the stream after layer l, the same point layer_<ll>.f32 records" — that reads as an absolute model layer, which is right given target_layer_ids; worth one explicit word, because "layer 46" in a MoE model could also be read as the 46th routed layer, and that mistake would produce a perfectly well-formed tap of the wrong tensor.

None of these block. The piece is coherent, the batched and step() paths are both covered, and the bit-for-bit test is the kind of check that makes the next PR (scoring a draft pair) meaningful rather than plausible.

@bong-water-water-bong

Copy link
Copy Markdown
Collaborator Author

Main merged in at 6221a8c so this PR picks up the CI fix for #2211 (f3f2db6, a ci.yml change, which a re-run of the old job would not use). No conflict; the PR's own files are unchanged.

Checked on this head: cmake --build rc 0; the PR's tests pass; full ctest -j6 1 failed of 244-245, which is npu_dmabuf_import. It fails the same way on all three of my branches ("invalid program file", its kernel argument is empty in my configuration), so it is not this change. Copyright and ci_targets checks rc 0.

ready for review at 6221a8c

@github-actions

Copy link
Copy Markdown

Persistent review updated to latest commit 6221a8c

@bong-water-water-bong bong-water-water-bong left a comment

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

VERDICT: READY at 6221a8c — the placement rules are a header-only pure function with 14 arms that I built and ran here, and the refusals are the half that matters.

$ g++ -std=c++17 -Inpu/forward/include -o lt tests/layer_taps_test.cpp && ./lt
layer_taps: ok (14 checks)          rc 0

The rules, and the four refusals

inline bool layer_taps_usable(const int* layers, int n, int model_layers) {
    for (int i = 0; i < n; i++) {
        if (layers[i] < 0 || layers[i] >= model_layers) return false;
        if (i > 0 && layers[i] <= layers[i - 1]) return false;      // strictly ascending
    }
    return true;
}

The arms are one per way a tap list can be wrong, and four of them are refusals: "a layer equal to the layer count is outside the model" (the off-by-one), "a repeated layer is refused", "a descending list is refused", "a negative layer is refused" — plus "the last layer of the model is usable" as the boundary that stops a guard being written one too tight. That is the shape the review wants: the boundary on both sides of the same comparison.

The two arms that are about the sink, not the list

"a pass that is not kept leaves the sink at the size it found"
"a failed outer pass takes back the rows of the passes inside it"

Rollback is the property a tap sink actually needs: a pass that fails must not leave partial rows for a later pass to read as if they were complete. LayerTapPass carrying a base offset is what makes that checkable, and the arms pin both directions (not-kept leaves the size; failed outer takes back the inner rows). That is the same discipline as #1970's lost-update fix, one layer up.

Header-only, no device, no pack — which is why I could run it at all, and why the rules are reviewable without a ticket.

One ask

layer_tap_put() writes with memcpy at (m * taps + j) * hidden, and the arms check the layout ("a row is one position with its taps side by side, in list order") — but there is no arm for j out of range (j >= taps would write into the next row's block, i.e. a silent overwrite rather than a refusal, and layer_tap_index() returns -1 for an unlisted layer so the caller has to guard). One check that layer_tap_index returning -1 is not passed to layer_tap_put — or an assert inside it — closes the only unchecked write in the header.

@bong-water-water-bong bong-water-water-bong left a comment

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

VERDICT at 6221a8c: READY

Re-review of 6221a8c. Since fa8d642 (READY from me) it adds 58426ba3, which answers my non-blocking note there (a failed pass left a zero-filled row block in the sink), plus a merge of main. The commit also fixes a real bug I had not seen: the last-layer tap was gated on err_.empty(), and err_ is never cleared, so one refused call blanked the last layer of every later pass.

Code: step() and step_batch() are now thin wrappers around step_impl / step_batch_impl, each holding a LayerTapPass that rewinds the sink to its size at entry unless the pass returned true (and, for the batched pass, no tapped layer changed row count: tap_fault_). The bodies call each other: the per-token fallback in step_batch_impl calls step() at :3826/:3859, and step_batch_all/embeds call step_batch(). The guards therefore nest, and they nest correctly. An inner success keeps its one row. Any failure rewinds the outer pass to where the whole pass began.

Checked at this head (ONEBIT_NPU=ON ONEBIT_NPU_HOST=ON build):

  • layer_taps_test: ok (14 checks).
  • layer_taps_model_test on ~/models/npu/Qwen3-0.6B: PASS layer_taps_model: 15 checks on 28 layers, H 1024.
  • Mutant, with LayerTapPass's destructor never rewinding: layer_taps_test rc 1, 3 FAIL. The model test still passes it, since no pass fails there, so the unit arm is the guard. After restoring the header the tree is clean and the test passes again.
  • It merges cleanly onto main. CI build at 6221a8c is green: 100% tests passed, 0 tests failed out of 243.

@bong-water-water-bong
bong-water-water-bong merged commit bc4d191 into main Nov 11, 2026
10 checks passed
@bong-water-water-bong
bong-water-water-bong deleted the cl-npuperf/1780-layer-taps branch November 11, 2026 21:33
bong-water-water-bong pushed a commit that referenced this pull request Nov 11, 2026
…ght_keys and layer_taps (engine#1959)

The one conflict was the include block of npu/forward/src/q4nx_forward.cpp. This PR
adds batch_split.h; main added weight_keys.h (#2214) and layer_taps.h (#2180) at the
same place. All three are kept. The diff against main is unchanged: the PR's
five files, +279/-1.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant