From fa8d6422fd02c761148473b871757781a15da3f1 Mon Sep 17 00:00:00 2001 From: bong-water-water-bong Date: Wed, 11 Nov 2026 08:37:22 -0400 Subject: [PATCH 1/2] npu: set_layer_taps(), the target's hidden states at chosen layers for 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 --- CMakeLists.txt | 13 +++ docs/npu.md | 6 ++ npu/forward/include/layer_taps.h | 53 +++++++++++ npu/forward/include/q4nx_forward.h | 12 +++ npu/forward/src/q4nx_forward.cpp | 35 ++++++++ tests/layer_taps_model_test.cpp | 139 +++++++++++++++++++++++++++++ tests/layer_taps_test.cpp | 64 +++++++++++++ 7 files changed, 322 insertions(+) create mode 100644 npu/forward/include/layer_taps.h create mode 100644 tests/layer_taps_model_test.cpp create mode 100644 tests/layer_taps_test.cpp diff --git a/CMakeLists.txt b/CMakeLists.txt index d6823978d..fbf6665ad 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -1583,6 +1583,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) @@ -1880,6 +1885,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 @@ -1910,6 +1922,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 diff --git a/docs/npu.md b/docs/npu.md index c17e721f8..d0c81799d 100644 --- a/docs/npu.md +++ b/docs/npu.md @@ -2667,6 +2667,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_.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_.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")`, diff --git a/npu/forward/include/layer_taps.h b/npu/forward/include/layer_taps.h new file mode 100644 index 000000000..7e079f3a1 --- /dev/null +++ b/npu/forward/include/layer_taps.h @@ -0,0 +1,53 @@ +// 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`. + +#pragma once + +#include +#include + +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)); +} + +} // namespace onebit::npu diff --git a/npu/forward/include/q4nx_forward.h b/npu/forward/include/q4nx_forward.h index 9618e07f4..84599eeb8 100644 --- a/npu/forward/include/q4nx_forward.h +++ b/npu/forward/include/q4nx_forward.h @@ -221,6 +221,14 @@ 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& 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. + bool set_layer_taps(const std::vector& layers, std::vector* 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& out) const { @@ -240,6 +248,10 @@ class Q4nxNpuForward { bool load_norms(); std::vector last_hidden_; // last_hidden(): the post-final-norm state of the last position std::vector last_residual_; // and one norm earlier: the raw residual of the last position + std::vector tap_layers_; // engine#1780: set_layer_taps() + std::vector* tap_sink_ = nullptr; + size_t tap_base_ = 0; // where the current pass's rows start in *tap_sink_ + void tap_layer(int layer, const float* X, int M); std::vector 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. diff --git a/npu/forward/src/q4nx_forward.cpp b/npu/forward/src/q4nx_forward.cpp index 37696df84..86abd43f6 100644 --- a/npu/forward/src/q4nx_forward.cpp +++ b/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" @@ -3411,6 +3412,34 @@ 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& layers, std::vector* 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) return; // a pass that changed width mid-way + 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 @@ -3679,6 +3708,9 @@ bool Q4nxNpuForward::step_batch(const std::vector& tokens, int start_pos, std::vector 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; @@ -4074,6 +4106,7 @@ bool Q4nxNpuForward::step_batch(const std::vector& tokens, int start_pos, // residual here produced garbage logits (the M=1 self-check caught it). std::vector fnorm(H); bool ok = true; + if (tap_sink_ && err_.empty()) 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 @@ -4860,6 +4893,7 @@ bool Q4nxNpuForward::step(int token, int pos, std::vector& 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 @@ -5182,6 +5216,7 @@ bool Q4nxNpuForward::step(int token, int pos, std::vector& logits, int ca } if (ph) { g_ph.layer = ph_ms(tp_lap); tp_lap = std::chrono::steady_clock::now(); } + if (tap_sink_ && err_.empty()) 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 diff --git a/tests/layer_taps_model_test.cpp b/tests/layer_taps_model_test.cpp new file mode 100644 index 000000000..10a290c27 --- /dev/null +++ b/tests/layer_taps_model_test.cpp @@ -0,0 +1,139 @@ +// 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_model_test.cpp -- Q4nxNpuForward::set_layer_taps() on a real container (engine#1780). +// +// Host-GEMM, no device; skips (77) unless ONEBIT_TEST_Q4NX names a model directory. What it holds the hook to: +// - a batched pass of M tokens appends M rows of [taps * H]; +// - the last layer's block of the last row is, bit for bit, the residual the forward itself keeps +// (last_residual()), so the tap is the stream and not a copy taken at the wrong point; +// - the same holds for per-token step() passes, and each appends exactly one row; +// - layer 0's block is, bit for bit, what NPU_INFER_DUMP_DIR writes for layer 0 in the same pass -- an instrument +// that already existed and writes at the end of the layer body, so a tap that is one layer off cannot agree; +// - a step() pass at position 0 gives the same row as the batched pass gave for that position, closely (the two +// paths are not bit-identical in general), which holds step()'s tap to the same layers; +// - a list that is not usable is refused and leaves the previous taps in place; an empty list turns them off. + +#include +#include +#include +#include +#include +#include +#include +#include + +#include "q4nx_forward.h" + +static int g_fail = 0, g_n = 0; +static void check(bool ok, const char* what) { + ++g_n; + if (!ok) { std::printf("FAIL %s\n", what); ++g_fail; } +} + +int main() { + const char* dir = std::getenv("ONEBIT_TEST_Q4NX"); + if (!dir || !*dir) { + std::printf("SKIP: set ONEBIT_TEST_Q4NX to a Q4NX model directory\n"); + return 77; + } + setenv("NPU_INFER_HOST_GEMM", "1", 1); + Q4nxNpuForward fwd; + const std::string model = std::string(dir) + "/model.q4nx"; + if (!fwd.init_host(model.c_str(), std::string(dir) + "/kernels")) { + std::printf("SKIP: init_host failed for %s (%s)\n", model.c_str(), fwd.last_error()); + return 77; + } + std::vector one; + fwd.embed_row(1, one); + const size_t H = one.size(); + const int NL = fwd.layer_count(); + if (NL < 2) { std::printf("SKIP: a %d-layer model has no two layers to tap\n", NL); return 77; } + const std::vector layers = {0, NL - 1}; + std::vector sink, logits; + + check(!fwd.set_layer_taps({NL - 1, 0}, &sink), "a descending list is refused"); + check(!fwd.set_layer_taps({0, NL}, &sink), "a layer outside the model is refused"); + check(fwd.set_layer_taps(layers, &sink), "the first and the last layer are accepted"); + + // The existing per-layer dump, into a directory beside the test's working directory. + namespace fs = std::filesystem; + const fs::path dump = fs::current_path() / ("layer_taps_dump." + std::to_string(getpid())); + fs::create_directories(dump); + setenv("NPU_INFER_DUMP_DIR", dump.c_str(), 1); + + const std::vector tokens = {1, 2, 3}; + const bool batched = !fwd.noncausal() ? fwd.step_batch_all(tokens, 0, logits) : fwd.step_batch(tokens, 0, logits); + check(batched, "the batched pass runs"); + check(sink.size() == tokens.size() * 2 * H, "a batched pass of 3 tokens appends 3 rows of 2 x H"); + if (batched && sink.size() == tokens.size() * 2 * H) { + const std::vector& res = fwd.last_residual(); + const float* last = sink.data() + (2 * 2 + 1) * H; // row 2, tap 1 + check(res.size() == H && std::memcmp(last, res.data(), H * sizeof(float)) == 0, + "batched: the last layer's block of the last row is the forward's own residual, bit for bit"); + check(std::memcmp(sink.data() + (2 * 2 + 0) * H, last, H * sizeof(float)) != 0, + "batched: layer 0's block is not the last layer's"); + std::vector ref(tokens.size() * H); + FILE* f = std::fopen((dump / "layer_00.f32").c_str(), "rb"); + const size_t got = f ? std::fread(ref.data(), sizeof(float), ref.size(), f) : 0; + if (f) std::fclose(f); + bool same = got == ref.size(); + for (size_t m = 0; m < tokens.size() && same; m++) + same = std::memcmp(sink.data() + (m * 2 + 0) * H, ref.data() + m * H, H * sizeof(float)) == 0; + check(same, "batched: layer 0's block of every row is what NPU_INFER_DUMP_DIR wrote for layer 0, bit for bit"); + } + unsetenv("NPU_INFER_DUMP_DIR"); + std::error_code ec; + fs::remove_all(dump, ec); + const std::vector batch_rows = sink; + + if (!fwd.noncausal()) { + sink.clear(); + fwd.reset(); + bool ok = true; + size_t grew = 0; + for (int i = 0; i < 3 && ok; i++) { + const size_t before = sink.size(); + ok = fwd.step(tokens[size_t(i)], i, logits); + grew = sink.size() - before; + ok = ok && grew == 2 * H; + } + check(ok, "each step() appends exactly one row of 2 x H"); + if (ok && batch_rows.size() >= 2 * H) { // position 0: the same token at the same place in both paths + double peak = 0.0, worst = 0.0; + for (size_t i = 0; i < 2 * H; i++) { + peak = std::max(peak, std::fabs(double(batch_rows[i]))); + worst = std::max(worst, std::fabs(double(batch_rows[i]) - double(sink[i]))); + } + std::printf("step() vs batched, position 0, both taps: max difference %.3g against a peak of %.3g\n", worst, peak); + check(worst <= 1e-3 * peak, "step(): position 0's row is the batched pass's row for it, to 1e-3 of the peak"); + } + if (ok) { + const std::vector& res = fwd.last_residual(); + check(res.size() == H && std::memcmp(sink.data() + (2 * 2 + 1) * H, res.data(), H * sizeof(float)) == 0, + "step(): the last layer's block of the last row is the forward's own residual, bit for bit"); + } + } + + const size_t held = sink.size(); + check(fwd.set_layer_taps({}, nullptr), "an empty list turns the taps off"); + fwd.reset(); + const bool again = fwd.noncausal() ? fwd.step_batch(tokens, 0, logits) : fwd.step(1, 0, logits); + check(again && sink.size() == held, "with the taps off a pass appends nothing"); + + if (g_fail) { std::printf("layer_taps_model: %d of %d checks FAILED\n", g_fail, g_n); return 1; } + std::printf("PASS layer_taps_model: %d checks on %d layers, H %zu\n", g_n, NL, H); + return 0; +} diff --git a/tests/layer_taps_test.cpp b/tests/layer_taps_test.cpp new file mode 100644 index 000000000..7fa75a604 --- /dev/null +++ b/tests/layer_taps_test.cpp @@ -0,0 +1,64 @@ +// 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_test.cpp -- the three rules of npu/forward/include/layer_taps.h (engine#1780), without a model. + +#include +#include + +#include "layer_taps.h" + +using namespace onebit::npu; + +static int g_fail = 0, g_n = 0; +static void check(bool ok, const char* what) { + ++g_n; + if (!ok) { std::printf("FAIL %s\n", what); ++g_fail; } +} + +int main() { + { + const int a[] = {5, 11, 17, 23, 35}; + check(layer_taps_usable(a, 5, 42), "the Bolt draft's five layers are usable on a 42-layer model"); + check(!layer_taps_usable(a, 5, 35), "a layer equal to the layer count is outside the model"); + check(layer_taps_usable(a, 5, 36), "the last layer of the model is usable"); + const int d[] = {5, 5}, r[] = {11, 5}, n[] = {-1, 3}; + check(!layer_taps_usable(d, 2, 42), "a repeated layer is refused"); + check(!layer_taps_usable(r, 2, 42), "a descending list is refused"); + check(!layer_taps_usable(n, 2, 42), "a negative layer is refused"); + check(!layer_taps_usable(a, 0, 42) && !layer_taps_usable(nullptr, 5, 42) && !layer_taps_usable(a, 5, 0), + "an empty list, a null list and a model without layers are refused"); + check(layer_tap_index(a, 5, 5) == 0 && layer_tap_index(a, 5, 23) == 3 && layer_tap_index(a, 5, 35) == 4, + "a listed layer owns the column block of its place in the list"); + check(layer_tap_index(a, 5, 6) == -1 && layer_tap_index(a, 5, 0) == -1, "an unlisted layer owns none"); + } + { // 2 positions, 3 taps, 2 wide: position m, tap j lands at (m * 3 + j) * 2, whatever order the layers arrive in + std::vector rows(2 * 3 * 2, -1.0f); + for (int j = 2; j >= 0; j--) // layer-major, last tap first, as a batched pass would not -- order-free + for (int m = 0; m < 2; m++) { + const float x[2] = {float(100 * m + 10 * j), float(100 * m + 10 * j + 1)}; + layer_tap_put(rows.data(), m, j, 3, 2, x); + } + bool ok = true; + for (int m = 0; m < 2; m++) + for (int j = 0; j < 3; j++) + for (int k = 0; k < 2; k++) + ok = ok && rows[size_t((m * 3 + j) * 2 + k)] == float(100 * m + 10 * j + k); + check(ok, "a row is one position with its taps side by side, in list order"); + } + if (g_fail) { std::printf("layer_taps: %d of %d checks FAILED\n", g_fail, g_n); return 1; } + std::printf("layer_taps: ok (%d checks)\n", g_n); + return 0; +} From 58426ba38a48793fa36099244605b23f2aecc33f Mon Sep 17 00:00:00 2001 From: bong-water-water-bong Date: Wed, 11 Nov 2026 09:56:51 -0400 Subject: [PATCH 2/2] layer taps: a failed pass appends nothing, and an earlier refusal no 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 --- npu/forward/include/layer_taps.h | 17 +++++++++++++++ npu/forward/include/q4nx_forward.h | 9 ++++++++ npu/forward/src/q4nx_forward.cpp | 35 +++++++++++++++++++++++++----- tests/layer_taps_model_test.cpp | 15 ++++++++++++- tests/layer_taps_test.cpp | 25 +++++++++++++++++++++ 5 files changed, 94 insertions(+), 7 deletions(-) diff --git a/npu/forward/include/layer_taps.h b/npu/forward/include/layer_taps.h index 7e079f3a1..ec93f7587 100644 --- a/npu/forward/include/layer_taps.h +++ b/npu/forward/include/layer_taps.h @@ -23,11 +23,16 @@ // - 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 #include +#include namespace onebit::npu { @@ -50,4 +55,16 @@ inline void layer_tap_put(float* rows, int m, int j, int taps, int hidden, const 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* sink; + size_t base; + bool keep = false; + explicit LayerTapPass(std::vector* 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 diff --git a/npu/forward/include/q4nx_forward.h b/npu/forward/include/q4nx_forward.h index 84599eeb8..9c501b687 100644 --- a/npu/forward/include/q4nx_forward.h +++ b/npu/forward/include/q4nx_forward.h @@ -227,6 +227,10 @@ class Q4nxNpuForward { // 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& layers, std::vector* 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` @@ -251,7 +255,12 @@ class Q4nxNpuForward { std::vector tap_layers_; // engine#1780: set_layer_taps() std::vector* 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& logits, int cache_pos); + bool step_batch_impl(const std::vector& tokens, int start_pos, std::vector& logits); std::vector 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. diff --git a/npu/forward/src/q4nx_forward.cpp b/npu/forward/src/q4nx_forward.cpp index 86abd43f6..4b3171964 100644 --- a/npu/forward/src/q4nx_forward.cpp +++ b/npu/forward/src/q4nx_forward.cpp @@ -3435,7 +3435,10 @@ void Q4nxNpuForward::tap_layer(int layer, const float* X, int M) { 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) return; // a pass that changed width mid-way + 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); } @@ -3540,8 +3543,28 @@ bool Q4nxNpuForward::step_batch_all(const std::vector& tokens, int start_po return ok; } -bool Q4nxNpuForward::step_batch(const std::vector& tokens, int start_pos, - std::vector& 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& tokens, int start_pos, std::vector& 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& 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& tokens, int start_pos, + std::vector& 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. @@ -4106,7 +4129,7 @@ bool Q4nxNpuForward::step_batch(const std::vector& tokens, int start_pos, // residual here produced garbage logits (the M=1 self-check caught it). std::vector fnorm(H); bool ok = true; - if (tap_sink_ && err_.empty()) tap_layer(c_.NL - 1, X.data(), M); + 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 @@ -4836,7 +4859,7 @@ bool Q4nxNpuForward::mlp(const float* gu, float* h, int im) { // -------------------------------------------------------------------------- // step // -------------------------------------------------------------------------- -bool Q4nxNpuForward::step(int token, int pos, std::vector& logits, int cache_pos) { +bool Q4nxNpuForward::step_impl(int token, int pos, std::vector& 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 @@ -5216,7 +5239,7 @@ bool Q4nxNpuForward::step(int token, int pos, std::vector& logits, int ca } if (ph) { g_ph.layer = ph_ms(tp_lap); tp_lap = std::chrono::steady_clock::now(); } - if (tap_sink_ && err_.empty()) tap_layer(c_.NL - 1, x.data(), 1); + 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 diff --git a/tests/layer_taps_model_test.cpp b/tests/layer_taps_model_test.cpp index 10a290c27..702deaeef 100644 --- a/tests/layer_taps_model_test.cpp +++ b/tests/layer_taps_model_test.cpp @@ -24,7 +24,10 @@ // that already existed and writes at the end of the layer body, so a tap that is one layer off cannot agree; // - a step() pass at position 0 gives the same row as the batched pass gave for that position, closely (the two // paths are not bit-identical in general), which holds step()'s tap to the same layers; -// - a list that is not usable is refused and leaves the previous taps in place; an empty list turns them off. +// - a list that is not usable is refused and leaves the previous taps in place; an empty list turns them off; +// - a call that was refused (a token outside the vocabulary) appends nothing, and the passes AFTER it still carry +// the last layer's block. The forward never clears its error text, so anything that gated a tap on "no error so +// far" would hand every later pass zeros there (review of fa8d642). #include #include @@ -74,6 +77,16 @@ int main() { fs::create_directories(dump); setenv("NPU_INFER_DUMP_DIR", dump.c_str(), 1); + // A refused call first, so the forward carries an error text from here on. + { + const std::vector bad = {1, -1}; + const bool refused_batch = !fwd.step_batch(bad, 0, logits); + const bool refused_step = !fwd.step(-1, 0, logits); + check(refused_batch && refused_step && sink.empty(), "a refused call appends nothing"); + check(fwd.last_error()[0] != '\0', "and the forward now carries an error text, as it will for the rest of its life"); + fwd.reset(); + } + const std::vector tokens = {1, 2, 3}; const bool batched = !fwd.noncausal() ? fwd.step_batch_all(tokens, 0, logits) : fwd.step_batch(tokens, 0, logits); check(batched, "the batched pass runs"); diff --git a/tests/layer_taps_test.cpp b/tests/layer_taps_test.cpp index 7fa75a604..f81b1a51b 100644 --- a/tests/layer_taps_test.cpp +++ b/tests/layer_taps_test.cpp @@ -58,6 +58,31 @@ int main() { ok = ok && rows[size_t((m * 3 + j) * 2 + k)] == float(100 * m + 10 * j + k); check(ok, "a row is one position with its taps side by side, in list order"); } + { // a pass that is not kept takes its rows back; a kept one leaves them; passes nest + std::vector sink(4, 1.0f); + { + LayerTapPass pass(&sink); + sink.resize(10, 0.0f); + } + check(sink.size() == 4, "a pass that is not kept leaves the sink at the size it found"); + { + LayerTapPass pass(&sink); + sink.resize(10, 2.0f); + pass.keep = true; + } + check(sink.size() == 10 && sink[3] == 1.0f && sink[9] == 2.0f, "a kept pass leaves its rows, and the earlier ones"); + { + LayerTapPass outer(&sink); + { + LayerTapPass inner(&sink); + sink.resize(16, 3.0f); + inner.keep = true; // the inner pass succeeded ... + } + } // ... inside an outer one that did not + check(sink.size() == 10, "a failed outer pass takes back the rows of the passes inside it"); + { LayerTapPass none(nullptr); } + check(true, "a pass with no sink does nothing"); + } if (g_fail) { std::printf("layer_taps: %d of %d checks FAILED\n", g_fail, g_n); return 1; } std::printf("layer_taps: ok (%d checks)\n", g_n); return 0;