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

CMakeLists.txt

@@ -1610,6 +1610,11 @@ add_test(NAME batch_check_compare COMMAND batch_check_compare_test)
 # engine#1780: the DFlash draft block's acceptance arithmetic -- tokens = accepted + bonus, the break-even rule, and the
 # published OGA-route counts as known answers. Header-only and host-only: no container, no device, no NPU runtime, so it
 # runs wherever this file is built. Registered at top level with the same hook as the row above, for the same reason.
+# engine#1780: where a tapped layer's rows go (npu/forward/include/layer_taps.h), header-only, no model.
+add_executable(layer_taps_test tests/layer_taps_test.cpp)
+target_include_directories(layer_taps_test PRIVATE npu/forward/include)
+add_test(NAME layer_taps COMMAND layer_taps_test)
+add_dependencies(onebit layer_taps_test)
 add_executable(dflash_accept_test tests/dflash_accept_test.cpp)
 target_include_directories(dflash_accept_test PRIVATE npu/forward/include)
 add_test(NAME dflash_accept COMMAND dflash_accept_test)
@@ -1918,6 +1923,13 @@ endif()
     target_link_libraries(gdn_batch_order_test PRIVATE onebit_npu_forward)
     add_dependencies(npu_sampling_test gdn_batch_order_test)
     add_test(NAME gdn_batch_order COMMAND gdn_batch_order_test)
+    # engine#1780: set_layer_taps() on a real container -- the last layer's tap is the forward's own residual, bit for
+    # bit, in the batched pass and in step(). Host-GEMM, no device; skips (77) unless ONEBIT_TEST_Q4NX names a model.
+    add_executable(layer_taps_model_test tests/layer_taps_model_test.cpp)
+    target_include_directories(layer_taps_model_test PRIVATE npu/forward/include)
+    target_link_libraries(layer_taps_model_test PRIVATE onebit_npu_forward)
+    add_dependencies(npu_sampling_test layer_taps_model_test)
+    add_test(NAME layer_taps_model COMMAND layer_taps_model_test)
     # engine#1905: the NPU_BATCH_CHECK oracle restores the prefix it replays from instead of resetting it, so a chunk
     # at start_pos > 0 is compared (and then run) from the state the batched pass continues. Host-GEMM, no device;
     # skips (77) unless ONEBIT_TEST_Q4NX names a container. This is the dev-box row (the private runtime, a real
@@ -1948,6 +1960,7 @@ endif()
     endif()
     set_tests_properties(mtp_batch_all PROPERTIES SKIP_RETURN_CODE 77)
     set_tests_properties(gdn_batch_order PROPERTIES SKIP_RETURN_CODE 77)
+    set_tests_properties(layer_taps_model PROPERTIES SKIP_RETURN_CODE 77)
     set_tests_properties(batch_check_state PROPERTIES SKIP_RETURN_CODE 77)
     # The fast-lane end-to-end test (npu_lane_e2e) is retired: its kernels were FastFlowLM-derived
     # and its reference logits are lost. The NPU check is the private dx add-on's from-source

docs/npu.md

@@ -2861,6 +2861,12 @@ can fail a run.
 variable and never sees that line is running a binary built before the dump existed (#1977): rebuild it, rather than
 looking for a code path that skips the dump.
 
+A caller inside the process that needs the same stream as data, not files, uses
+`Q4nxNpuForward::set_layer_taps(layers, &sink)` (engine#1780): every pass then appends one row per position to the
+sink, `taps x H` wide, the listed layers side by side. It works for `step()` and for the batched passes, and it is what
+a DFlash draft's `target_layer_ids` are read with (docs/amd-drafters.md). Layer `l`'s block is the stream after layer
+`l`, the same point `layer_<ll>.f32` records; `layer_taps_model_test` holds the two to each other bit for bit.
+
 The files these dumps write (`NPU_INFER_DUMP_DIR`'s `layer_<ll>.f32`, `NPU_I8_GS_DUMP`'s weight-scale pairs) are
 created owner-only (0600) and never through a symlink at the final path component, whatever the umask, and a dump file
 an older build left at 0666 is narrowed to 0600 when it is rewritten; the head-image cache's files are at most 0644. Both go through `npu/forward/include/file_create.h`, so no write site uses `fopen(…, "w")`,

npu/forward/include/layer_taps.h

@@ -0,0 +1,70 @@
+// Copyright 2026 bong-water-water-bong
+// SPDX-License-Identifier: Apache-2.0
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+// layer_taps.h -- where a tapped layer's rows go (engine#1780).
+//
+// A DFlash draft reads the target model's hidden states at a few layers (its target_layer_ids). The forward collects
+// them with Q4nxNpuForward::set_layer_taps(); the three rules that decide what lands where are here, header-only, so
+// a test can hold them without a model:
+//
+//   - which lists are usable: ascending, no repeats, every layer inside the model;
+//   - which column block a layer owns: its index in the list, or none;
+//   - where a row goes: position m of a pass, tap j of T, H wide -> rows[(m * T + j) * H ...]. That is one row per
+//     position with the tapped layers side by side, the layout dflash_forward.h takes as `taps`.
+//
+// And one about failure: LayerTapPass takes the sink back to the size it had when the pass began unless the pass is
+// marked kept. A pass opens its rows when the first listed layer arrives, so a pass that fails later would otherwise
+// leave rows whose remaining blocks are zeros, and zeros read as data.
+
+#pragma once
+
+#include <cstddef>
+#include <cstring>
+#include <vector>
+
+namespace onebit::npu {
+
+inline bool layer_taps_usable(const int* layers, int n, int model_layers) {
+    if (!layers || n <= 0 || model_layers <= 0) return false;
+    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;
+    }
+    return true;
+}
+
+inline int layer_tap_index(const int* layers, int n, int layer) {
+    for (int i = 0; i < n; i++)
+        if (layers[i] == layer) return i;
+    return -1;
+}
+
+inline void layer_tap_put(float* rows, int m, int j, int taps, int hidden, const float* x) {
+    std::memcpy(rows + (size_t(m) * size_t(taps) + size_t(j)) * size_t(hidden), x, size_t(hidden) * sizeof(float));
+}
+
+struct LayerTapPass {
+    std::vector<float>* sink;
+    size_t base;
+    bool keep = false;
+    explicit LayerTapPass(std::vector<float>* s) : sink(s), base(s ? s->size() : 0) {}
+    ~LayerTapPass() {
+        if (sink && !keep && sink->size() > base) sink->resize(base);
+    }
+    LayerTapPass(const LayerTapPass&) = delete;
+    LayerTapPass& operator=(const LayerTapPass&) = delete;
+};
+
+}  // namespace onebit::npu

npu/forward/include/q4nx_forward.h

@@ -221,6 +221,18 @@ class Q4nxNpuForward {
     // The post-final-norm hidden of EVERY row of the last batched pass (M x H, row-major): a speculative drafter
     // has to advance its own stream over the rows that were actually accepted, and those are exactly these.
     const std::vector<float>& last_hidden_rows() const { return last_hidden_rows_; }
+    // engine#1780: the residual stream after the listed layers, for a draft that reads the target's hidden states
+    // (a DFlash draft's target_layer_ids). `layers` must be ascending and inside the model; every pass then appends
+    // one row per position to *sink, [taps * H] wide, the listed layers side by side in the order given -- the
+    // layout dflash_forward.h's `taps` argument takes. Layer l's row is the stream as layer l + 1 receives it (after
+    // layer l's own closing scale). Works for step() and for the batched passes. An empty list or a null sink turns
+    // it off; the sink must outlive the passes. Returns false, changing nothing, on a list that is not usable.
+    // A pass that returns false appends nothing: whatever it had written is taken back, so the sink only ever holds
+    // rows of passes that succeeded. A call refused earlier on the same forward does not affect later passes.
+    // Not for use under NPU_BATCH_CHECK: that oracle replays positions through step(), and the replays would be
+    // appended as if they were new positions.
+    bool set_layer_taps(const std::vector<int>& layers, std::vector<float>* sink);
+    int layer_count() const { return c_.NL; }   // layers of the loaded model; 0 before init
     // One row of the embedding table (what a token means to the model): the MTP drafter's `pre_fc_norm_embedding`
     // input.  It reads the same table the forward uses, so the drafter cannot drift from the model's embeddings.
     void embed_row(int token, std::vector<float>& out) const {
@@ -240,6 +252,15 @@ class Q4nxNpuForward {
     bool load_norms();
     std::vector<float> last_hidden_;  // last_hidden(): the post-final-norm state of the last position
     std::vector<float> last_residual_;  // and one norm earlier: the raw residual of the last position
+    std::vector<int>    tap_layers_;           // engine#1780: set_layer_taps()
+    std::vector<float>* tap_sink_ = nullptr;
+    size_t              tap_base_ = 0;         // where the current pass's rows start in *tap_sink_
+    bool                tap_fault_ = false;    // a tapped layer arrived with a different row count than the pass opened with
+    void tap_layer(int layer, const float* X, int M);
+    // The bodies of step() and step_batch(). The public functions wrap them so that a failed pass leaves the tap
+    // sink as it found it (layer_taps.h, LayerTapPass).
+    bool step_impl(int token, int pos, std::vector<float>& logits, int cache_pos);
+    bool step_batch_impl(const std::vector<int>& tokens, int start_pos, std::vector<float>& logits);
     std::vector<float> last_hidden_rows_;  // [M][H] normed hidden of the last batched pass
     // Set only while step_batch_all() runs: the batched epilogue then norms every row and writes M x NV
     // logits there instead of only the last row's.  Null on every other path, so nothing else changes.

npu/forward/src/q4nx_forward.cpp

@@ -24,6 +24,7 @@
 // Numerics follow tools/q4nx_forward_ref.py so the two can be compared
 // element-wise.
 #include "q4nx_forward.h"
+#include "layer_taps.h"      // engine#1780: where a tapped layer's rows go in the sink
 #include "pad_head_rows.h"   // engine#1990: the padded-head row rule, header-only for its arm
 #include "batch_width.h"
 #include "ffn_cache_policy.h"
@@ -3613,6 +3614,37 @@ void Q4nxNpuForward::npu_gemm_rows(int role, const float* A, int M, int K, int N
     op_line(sc, cur_layer_, M, out, N);
 }
 
+bool Q4nxNpuForward::set_layer_taps(const std::vector<int>& layers, std::vector<float>* sink) {
+    if (layers.empty() || !sink) {
+        tap_layers_.clear();
+        tap_sink_ = nullptr;
+        return true;
+    }
+    if (!onebit::npu::layer_taps_usable(layers.data(), int(layers.size()), c_.NL)) return false;
+    tap_layers_ = layers;
+    tap_sink_ = sink;
+    tap_base_ = sink->size();
+    return true;
+}
+
+// One tapped layer of one pass: M rows of H, written into the pass's rows of the sink at this layer's column block.
+void Q4nxNpuForward::tap_layer(int layer, const float* X, int M) {
+    const int T = int(tap_layers_.size());
+    const int j = onebit::npu::layer_tap_index(tap_layers_.data(), T, layer);
+    if (j < 0 || !tap_sink_ || M <= 0) return;
+    const size_t H = size_t(c_.H);
+    if (j == 0) {   // the first listed layer opens the pass's rows
+        tap_base_ = tap_sink_->size();
+        tap_sink_->resize(tap_base_ + size_t(M) * size_t(T) * H);
+    }
+    if (tap_sink_->size() != tap_base_ + size_t(M) * size_t(T) * H) {   // not the rows this pass opened with
+        tap_fault_ = true;
+        return;
+    }
+    for (int m = 0; m < M; m++)
+        onebit::npu::layer_tap_put(tap_sink_->data() + tap_base_, m, j, T, int(H), X + size_t(m) * H);
+}
+
 // The folded container's sign table: magic "HAD1", u32 count, count x (u32 width, u32 offset),
 // then the f32 signs those offsets index (the GGUF's prism.hadamard.sign_widths/sign_values).
 // The plain-text overload. It is declared in q4nx_forward.h and used by npu/forward/embeds_equiv_probe.cpp, but only the
@@ -3713,8 +3745,28 @@ bool Q4nxNpuForward::step_batch_all(const std::vector<int>& tokens, int start_po
     return ok;
 }
 
-bool Q4nxNpuForward::step_batch(const std::vector<int>& tokens, int start_pos,
-                                std::vector<float>& logits) {
+// engine#1780: a pass that fails must not leave rows in the tap sink. The bodies below return false from many
+// places, so the rule is applied once, here, around each of them.
+bool Q4nxNpuForward::step_batch(const std::vector<int>& tokens, int start_pos, std::vector<float>& logits) {
+    onebit::npu::LayerTapPass pass(tap_sink_);
+    tap_fault_ = false;
+    const bool ok = step_batch_impl(tokens, start_pos, logits);
+    if (tap_fault_ && tap_sink_)
+        std::fprintf(stderr, "  [taps] a batched pass at position %d changed its row count between tapped layers; "
+                             "its rows are dropped\n", start_pos);
+    pass.keep = ok && !tap_fault_;
+    return ok;
+}
+
+bool Q4nxNpuForward::step(int token, int pos, std::vector<float>& logits, int cache_pos) {
+    onebit::npu::LayerTapPass pass(tap_sink_);
+    const bool ok = step_impl(token, pos, logits, cache_pos);
+    pass.keep = ok;
+    return ok;
+}
+
+bool Q4nxNpuForward::step_batch_impl(const std::vector<int>& tokens, int start_pos,
+                                     std::vector<float>& logits) {
     if (!ready_) { err_ = "not initialised"; return false; }
     // A noncausal pack reads the whole prompt at once, so a chunk that starts above 0 was already
     // encoded causally by an earlier call and cannot be repaired here.  Named, not silently causal.
@@ -3881,6 +3933,9 @@ bool Q4nxNpuForward::step_batch(const std::vector<int>& tokens, int start_pos,
 
     std::vector<float> QKV, CTX, O, ACT, HH, D, qkv_row, ctx_row;
     for (int l = 0; l < c_.NL; l++) {
+        // engine#1780: the stream as this layer receives it is the previous layer's output. Taken here and not at
+        // the end of the body, because the body leaves through more than one `continue`.
+        if (l > 0 && tap_sink_) tap_layer(l - 1, X.data(), M);
         cur_layer_ = l;   // the layer index the per-layer helpers read
         const bool cached = weight_cached(0);
         const int lt0 = (l < sem_.n_layer_types) ? sem_.layer_types[l] : Q4NX_LAYER_FULL;
@@ -4278,6 +4333,7 @@ bool Q4nxNpuForward::step_batch(const std::vector<int>& tokens, int start_pos,
     // residual here produced garbage logits (the M=1 self-check caught it).
     std::vector<float> fnorm(H);
     bool ok = true;
+    if (tap_sink_) tap_layer(c_.NL - 1, X.data(), M);
     last_residual_.assign(X.begin() + (size_t)(M - 1) * H, X.end());
     if (batch_all_out_) {
         // step_batch_all(): the verifier needs the model's own choice at every drafted position, so every
@@ -5014,7 +5070,7 @@ bool Q4nxNpuForward::mlp(const float* gu, float* h, int im) {
 // --------------------------------------------------------------------------
 // step
 // --------------------------------------------------------------------------
-bool Q4nxNpuForward::step(int token, int pos, std::vector<float>& logits, int cache_pos) {
+bool Q4nxNpuForward::step_impl(int token, int pos, std::vector<float>& logits, int cache_pos) {
     if (!ready_) { err_ = "not initialised"; return false; }
     // Generation is causal by construction (one position at a time, keys only in the past), so a pack
     // whose full-attention layers read the whole prompt cannot be run this way.  Refused by name
@@ -5071,6 +5127,7 @@ bool Q4nxNpuForward::step(int token, int pos, std::vector<float>& logits, int ca
     if (ph) { g_ph.embed = ph_ms(tp_step); tp_lap = std::chrono::steady_clock::now(); }
 
     for (int l = 0; l < c_.NL; l++) {
+        if (l > 0 && tap_sink_) tap_layer(l - 1, x.data(), 1);   // engine#1780, as in the batched body
         cur_layer_ = l;
         // Dense models: once this layer's projections are packed (after the
         // first token), the f32 weights are no longer read, so skip the per-token
@@ -5395,6 +5452,7 @@ bool Q4nxNpuForward::step(int token, int pos, std::vector<float>& logits, int ca
     }
     if (ph) { g_ph.layer = ph_ms(tp_lap); tp_lap = std::chrono::steady_clock::now(); }
 
+    if (tap_sink_) tap_layer(c_.NL - 1, x.data(), 1);
     last_residual_ = x;                        // one norm earlier, for a drafter that wants the residual
     rmsnorm(x.data(), final_norm_.data(), H, xn.data());
     last_hidden_ = xn;  // for last_hidden(): the post-final-norm state of this position