npu: the batch-check oracle compares every row, bitwise, instead of one argmax at the last position

CMakeLists.txt

@@ -353,6 +353,8 @@ target_include_directories(serve_watchdog_test PRIVATE app)
 add_test(NAME serve_watchdog COMMAND serve_watchdog_test)
 # The NPU route's sampling controls and GBNF output constraint (app/npu_sampling.h,
 # app/npu_grammar.h, docs/npu.md "Sampling and structured output"); header only, no device.
+# engine#1781: the comparator behind NPU_BATCH_CHECK -- pure arithmetic over two float arrays, so it runs with no model,
+# no device and no runtime. It is a header, so the test includes it directly and needs no link libraries.
 add_executable(npu_sampling_test tests/npu_sampling_test.cpp)
 target_include_directories(npu_sampling_test PRIVATE app)
 target_link_libraries(npu_sampling_test PRIVATE nlohmann_json::nlohmann_json)
@@ -814,6 +816,8 @@ add_test(NAME route_port_lock COMMAND ${Python3_EXECUTABLE} ${CMAKE_SOURCE_DIR}/
     # the HRX hardware gate's judges (docs/hw-gate.md) against a fake server and synthetic logs;
     # the gate itself needs the GPU and runs on Strix Halo (.github/workflows/hw-gate.yml)
     add_test(NAME hw_gate_judges COMMAND ${Python3_EXECUTABLE} ${CMAKE_SOURCE_DIR}/tests/hw_gate_test.py)
+    # engine#1781: the comparator behind NPU_BATCH_CHECK, host-only (no model, no device, no runtime):
+    # the batched verify pass is compared row for row against the sequential one, bitwise.
     add_test(NAME q4nx_pb16i COMMAND ${Python3_EXECUTABLE} ${CMAKE_SOURCE_DIR}/tests/q4nx_pb16i_test.py)
     # The padded-grid read (engine#1097/#1110): tools/q4nx_diff.py decode reads a container whose weight sits on a
     # WIDER tile grid than the model's dimensions (the FastFlowLM repack pads intermediate_size to 512 columns, IM

npu/forward/include/batch_check_compare.h

@@ -0,0 +1,87 @@
+// 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.
+//
+// batch_check_compare.h -- the comparator behind NPU_BATCH_CHECK (engine#1781).
+//
+// The oracle that guards a batched verify pass used to compare ONE argmax at ONE position: the last row's. That
+// answers "did the final argmax survive batching", which is not the question a speculative decoder asks -- it needs
+// every drafted position's LOGITS (the accept rule consumes the target's probability at the drafted token, and two
+// passes can agree on an argmax while disagreeing on the distribution around it).
+//
+// Pure and header-only on purpose: the comparison is arithmetic over two float arrays, so it can be pinned by a host
+// test with no device, no model and no runtime, while the device-only half stays a single call in the forward.
+
+#ifndef ONEBIT_NPU_FORWARD_BATCH_CHECK_COMPARE_H
+#define ONEBIT_NPU_FORWARD_BATCH_CHECK_COMPARE_H
+
+#include <cmath>
+#include <cstdint>
+#include <cstring>
+
+namespace onebit {
+namespace npu {
+
+struct RowMismatch {
+    int row = -1;            // first row that differs, -1 when every row matched
+    int index = -1;          // column of the first difference within that row
+    uint32_t seq_bits = 0;   // the two values, as bits, because a float print of two near-equal values reads as equal
+    uint32_t bat_bits = 0;
+    bool nan_mismatch = false;  // one side NaN and the other not (bit comparison alone would call two NaNs equal)
+};
+
+// Compare `rows` x `cols` float rows BITWISE. Returns the number of differing rows (0 == identical).
+//
+// Bitwise, not tolerance-based: the gate is "the batched pass is the same arithmetic as the sequential one", so a
+// value that differs in the last mantissa bit is a divergence to report, not a rounding to forgive. NaN is handled
+// explicitly because the usual bit comparison would consider two NaN payloads equal and, more importantly, a NaN on
+// one side with a finite number on the other is exactly the failure this oracle exists to name.
+inline int compare_rows_bitwise(const float* seq, const float* bat, int rows, int cols, RowMismatch* first) {
+    if (!seq || !bat || rows <= 0 || cols <= 0) return rows > 0 ? rows : 0;
+    int differing = 0;
+    for (int r = 0; r < rows; ++r) {
+        const float* s = seq + (size_t)r * (size_t)cols;
+        const float* b = bat + (size_t)r * (size_t)cols;
+        for (int i = 0; i < cols; ++i) {
+            const bool s_nan = std::isnan(s[i]), b_nan = std::isnan(b[i]);
+            if (s_nan || b_nan) {
+                if (s_nan && b_nan) continue;   // both not-a-number: the same non-answer, not a divergence
+                ++differing;
+                if (first && first->row < 0) {
+                    first->row = r; first->index = i; first->nan_mismatch = true;
+                    std::memcpy(&first->seq_bits, &s[i], 4);
+                    std::memcpy(&first->bat_bits, &b[i], 4);
+                }
+                break;                          // one report per row is enough to know the row differs
+            }
+            uint32_t sb = 0, bb = 0;
+            std::memcpy(&sb, &s[i], 4);
+            std::memcpy(&bb, &b[i], 4);
+            if (sb != bb) {
+                ++differing;
+                if (first && first->row < 0) {
+                    first->row = r; first->index = i;
+                    first->seq_bits = sb; first->bat_bits = bb;
+                }
+                break;
+            }
+        }
+    }
+    return differing;
+}
+
+}  // namespace npu
+}  // namespace onebit
+
+#endif  // ONEBIT_NPU_FORWARD_BATCH_CHECK_COMPARE_H

npu/forward/src/q4nx_forward.cpp

@@ -33,6 +33,7 @@
 #include <chrono>
 #include <cmath>
 #include <cstdio>
+#include "batch_check_compare.h"
 #include "residual-dump.h"
 #include <cstdlib>
 #include <cstring>
@@ -3338,16 +3339,25 @@ bool Q4nxNpuForward::step_batch(const std::vector<int>& tokens, int start_pos,
     // device state -- localises a batched-vs-sequential divergence immediately.
     static const bool batch_check = getenv("NPU_BATCH_CHECK") != nullptr;
     int seq_argmax = -1;
+    // engine#1781: the old oracle compared ONE argmax at ONE position -- the last row's -- which is not the question a
+    // speculative decoder asks. Every drafted position's LOGITS matter (the accept rule consumes the target's
+    // probability at the drafted token), so the sequential rows are kept whole and compared row for row, bitwise.
+    // M x vocab floats: 4-7 MB on the models this runs, allocated only under the debug flag.
+    std::vector<float> seq_rows, bat_rows;
     // Not with injected rows: the oracle's own sequential pass runs step(tokens[m]) on the placeholder ids this
     // entry point supplies, so it would compare against a zero-token pass and report a divergence that is not
     // there. A debug tool that lies is worse than one that says "not applicable".
     if (batch_check && !inject_embeds_) {
         std::vector<float> seq;
-        for (int m = 0; m < M; m++)
+        if (c_.NV > 0) seq_rows.reserve((size_t)M * (size_t)c_.NV);
+        for (int m = 0; m < M; m++) {
             if (!step(tokens[m], start_pos + m, seq)) return false;
+            seq_rows.insert(seq_rows.end(), seq.begin(), seq.end());   // every row, not only the last
+            if (m == M - 1)
+                for (size_t i = 0; i < seq.size(); i++)
+                    if (seq_argmax < 0 || seq[i] > seq[(size_t)seq_argmax]) seq_argmax = (int)i;
+        }
         reset();
-        for (size_t i = 0; i < seq.size(); i++)
-            if (seq_argmax < 0 || seq[i] > seq[(size_t)seq_argmax]) seq_argmax = (int)i;
     }
 
     // Batch phase attribution: the same partition step() prints, for an M-row verify pass, so the
@@ -3759,6 +3769,8 @@ bool Q4nxNpuForward::step_batch(const std::vector<int>& tokens, int start_pos,
             ok = run_lm_head_rows(norm_all, M, *batch_all_out_);
         }
         logits.assign(batch_all_out_->begin() + (size_t)(M - 1) * (size_t)c_.NV, batch_all_out_->end());
+        // engine#1781: keep the batched rows so the oracle can compare all of them, not just the last one's argmax.
+        if (batch_check) bat_rows.assign(batch_all_out_->begin(), batch_all_out_->end());
     } else {
         std::vector<float> last(X.begin() + (size_t)(M - 1) * H, X.end());
         rmsnorm(last.data(), final_norm_.data(), H, fnorm.data());
@@ -3786,6 +3798,16 @@ bool Q4nxNpuForward::step_batch(const std::vector<int>& tokens, int start_pos,
             if (logits[i] > logits[(size_t)ba]) ba = (int)i;
         std::fprintf(stderr, "[batchcheck] M=%d start=%d seq_argmax=%d batch_argmax=%d %s\n",
                      M, start_pos, seq_argmax, ba, seq_argmax == ba ? "MATCH" : "MISMATCH");
+        // engine#1781: and the comparison the gate actually asks for -- every row, bitwise.
+        if (!seq_rows.empty() && bat_rows.size() == seq_rows.size() && c_.NV > 0) {
+            onebit::npu::RowMismatch mm;
+            const int differing = onebit::npu::compare_rows_bitwise(seq_rows.data(), bat_rows.data(), M, c_.NV, &mm);
+            std::fprintf(stderr, "[batchcheck] rows M=%d identical=%d differing=%d", M, M - differing, differing);
+            if (differing > 0)
+                std::fprintf(stderr, " first row=%d index=%d seq=0x%08x bat=0x%08x%s", mm.row, mm.index,
+                             (unsigned)mm.seq_bits, (unsigned)mm.bat_bits, mm.nan_mismatch ? " (NaN)" : "");
+            std::fprintf(stderr, "\n");
+        }
     } else if (getenv("NPU_BATCH_DBG")) {
         std::fprintf(stderr, "[batch] M=%d start=%d batched (%d layers, two_widths=%d)\n",
                      M, start_pos, c_.NL, (int)two_widths);

tests/batch_check_compare_test.cpp

@@ -0,0 +1,115 @@
+// 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.
+//
+// batch_check_compare_test.cpp -- pins the comparator behind NPU_BATCH_CHECK (engine#1781), so the oracle's answer is
+// arithmetic that a host test can check rather than a device run that has to be trusted. No model, no device, no
+// runtime: it is two float arrays and a count.
+
+#include <cmath>
+#include <cstdio>
+#include <cstring>
+#include <limits>
+#include <vector>
+
+#include "batch_check_compare.h"
+
+using onebit::npu::compare_rows_bitwise;
+using onebit::npu::RowMismatch;
+
+static int failures = 0;
+
+static void check(bool ok, const char* what, const char* detail = "") {
+    std::printf("%s %s%s%s\n", ok ? "ok  " : "FAIL", what, ok ? "" : " -- ", ok ? "" : detail);
+    if (!ok) ++failures;
+}
+
+static std::vector<float> rows(int nrows, int cols, float base) {
+    std::vector<float> v((size_t)nrows * (size_t)cols);
+    for (size_t i = 0; i < v.size(); ++i) v[i] = base + (float)(i % 97) * 0.011f;
+    return v;
+}
+
+int main() {
+    const int R = 7, C = 64;                       // the shape the gate asks about: k = 7 rows of logits
+
+    {   // 1. identical buffers match, row for row
+        auto a = rows(R, C, 0.25f);
+        auto b = a;
+        RowMismatch m;
+        check(compare_rows_bitwise(a.data(), b.data(), R, C, &m) == 0, "identical rows report no divergence");
+        check(m.row == -1 && m.index == -1, "and leave the mismatch untouched");
+    }
+    {   // 2. a difference in the LAST row is found -- the old oracle's only case, so it must still work
+        auto a = rows(R, C, 0.25f);
+        auto b = a;
+        b[(size_t)(R - 1) * C + 5] += 1e-7f;
+        RowMismatch m;
+        int n = compare_rows_bitwise(a.data(), b.data(), R, C, &m);
+        char d[160];
+        std::snprintf(d, sizeof d, "n=%d row=%d index=%d", n, m.row, m.index);
+        check(n == 1 && m.row == R - 1 && m.index == 5, "a difference in the last row is found", d);
+    }
+    {   // 3. ...and a difference in an EARLY row is found, which the old oracle could not see at all
+        auto a = rows(R, C, 0.25f);
+        auto b = a;
+        b[(size_t)2 * C + 11] += 1e-7f;            // row 2 of 7: a drafted position, not the last one
+        RowMismatch m;
+        int n = compare_rows_bitwise(a.data(), b.data(), R, C, &m);
+        char d[160];
+        std::snprintf(d, sizeof d, "n=%d row=%d index=%d", n, m.row, m.index);
+        check(n == 1 && m.row == 2 && m.index == 11, "a difference in an early row is found too", d);
+    }
+    {   // 4. bit-exactness is the standard: one ulp apart is a divergence, not a rounding
+        auto a = rows(R, C, 1.0f);
+        auto b = a;
+        b[3] = std::nextafter(b[3], 2.0f);
+        RowMismatch m;
+        int n = compare_rows_bitwise(a.data(), b.data(), R, C, &m);
+        check(n == 1 && m.row == 0 && m.index == 3, "a one-ulp difference is reported (the bar is bit-exact)", "no divergence reported");
+    }
+    {   // 5. every differing row is counted, so a run reports the size of the divergence and not just its presence
+        auto a = rows(R, C, 0.25f);
+        auto b = a;
+        for (int r = 0; r < R; ++r) b[(size_t)r * C + 1] += 1e-6f;
+        RowMismatch m;
+        int n = compare_rows_bitwise(a.data(), b.data(), R, C, &m);
+        char d[80];
+        std::snprintf(d, sizeof d, "n=%d (wanted %d)", n, R);
+        check(n == R, "all seven differing rows are counted", d);
+        check(m.row == 0, "and the first one is the first row");
+    }
+    {   // 6. NaN on one side only is a divergence; NaN on both is not (the same non-answer)
+        const float nan = std::numeric_limits<float>::quiet_NaN();
+        auto a = rows(2, 4, 0.5f);
+        auto b = a;
+        b[6] = nan;                                 // row 1, col 2: batched has NaN, sequential does not
+        RowMismatch m;
+        check(compare_rows_bitwise(a.data(), b.data(), 2, 4, &m) == 1 && m.row == 1 && m.index == 2 && m.nan_mismatch,
+              "a one-sided NaN is reported as a mismatch");
+        auto c = a;
+        c[6] = nan;
+        auto d2 = a;
+        d2[6] = nan;
+        check(compare_rows_bitwise(c.data(), d2.data(), 2, 4, nullptr) == 0, "and two NaNs in the same place agree");
+    }
+    {   // 7. degenerate shapes are refusals, not silent success
+        auto a = rows(2, 4, 0.5f);
+        check(compare_rows_bitwise(nullptr, a.data(), 2, 4, nullptr) == 2, "a null side counts every row as differing");
+        check(compare_rows_bitwise(a.data(), a.data(), 0, 4, nullptr) == 0, "zero rows compares as equal (nothing to check)");
+    }
+
+    std::printf("batch_check_compare: %d failure(s)\n", failures);
+    return failures ? 1 : 0;
+}