npu: step_batch_all() -- per-row logits from one batched pass, for speculative verification

CMakeLists.txt

@@ -712,6 +712,15 @@ if(ONEBIT_NPU)
     target_sources(onebit PRIVATE app/forward_serve.cpp)
     target_link_libraries(onebit PRIVATE onebit_npu_forward)
     set(ONEBIT_NPU_KERNELS "" CACHE PATH "Kernel directory (layer_ctx{1,2,17}.elf, lmhead.elf, layer.pdi)")
+    # step_batch_all(): per-row logits from ONE batched pass -- the contract speculative verification
+    # rests on (a draft token is accepted only when the model's own choice at that position matches it).
+    # Host-GEMM mode, no device; needs a Q4NX container, so it skips (77) unless ONEBIT_TEST_Q4NX names one.
+    add_executable(mtp_batch_all_test tests/mtp_batch_all_test.cpp)
+    target_include_directories(mtp_batch_all_test PRIVATE npu/forward/include)
+    target_link_libraries(mtp_batch_all_test PRIVATE onebit_npu_forward)
+    add_dependencies(npu_sampling_test mtp_batch_all_test)
+    add_test(NAME mtp_batch_all COMMAND mtp_batch_all_test)
+    set_tests_properties(mtp_batch_all 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
     # check, outside this repository (docs/npu.md, "Checking the NPU").

npu/forward/include/q4nx_forward.h

@@ -197,6 +197,15 @@ class Q4nxNpuForward {
     bool step_batch(const std::vector<int>& tokens, int start_pos,
                     std::vector<float>& logits);
 
+    // The same batched pass, but returning the logits of EVERY row (M x NV, row-major) instead of only
+    // the last one.  Speculative decoding needs them: a draft token is accepted only if the model's own
+    // choice AT THAT POSITION is that token, so the verifier has to see each row's logits, not just the
+    // tail of the batch.  The layer maths is the batched body step_batch already runs -- only the final
+    // norm and the LM head are applied per row, and the head is launched once for all M rows (the M=16
+    // twin streams its weights once for the batch, exactly like the layers).
+    bool step_batch_all(const std::vector<int>& tokens, int start_pos,
+                        std::vector<float>& logits_all);
+
     // The last position's hidden state after the final norm, H floats: what a d3 readout head pools
     // (`pooling: last` in a pack's decision_config.json, npu/decision.h). Filled by the last
     // successful step() or step_batch(), empty before that. This is the vector the LM head reads, not
@@ -214,6 +223,9 @@ class Q4nxNpuForward {
 private:
     bool load_norms();
     std::vector<float> last_hidden_;  // last_hidden(): the post-final-norm state of the last position
+    // 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.
+    std::vector<float>* batch_all_out_ = nullptr;
     bool init_impl(xrt::device* dev, const char* model_path,
                    const std::string& artifact_dir);
     bool host_gemm_ = false;
@@ -276,6 +288,9 @@ class Q4nxNpuForward {
     float rmsnorm(const float* x, const float* w, int n, float* out) const;
     float rmsnorm_noweight(const float* x, int n, float* out) const;
     bool run_lm_head(const std::vector<float>& h, std::vector<float>& logits);
+    // M rows of hidden state -> M x NV logits, row-major.  Same head, same packed tiles: the rows ride
+    // one launch (per-row activation scales, which the batched launcher and the dequantiser already carry).
+    bool run_lm_head_rows(const std::vector<float>& h_rows, int M, std::vector<float>& logits_all);
 
     Q4nxForwardConfig c_;
     ModelWeights* mw_ = nullptr;

npu/forward/src/q4nx_forward.cpp

@@ -2778,6 +2778,16 @@ void Q4nxNpuForward::npu_gemm_rows(I8Ctx& ctx, const float* A, int M, int K, int
     }
 }
 
+bool Q4nxNpuForward::step_batch_all(const std::vector<int>& tokens, int start_pos,
+                                    std::vector<float>& logits_all) {
+    std::vector<float> last;
+    std::vector<float>* saved = batch_all_out_;
+    batch_all_out_ = &logits_all;
+    const bool ok = step_batch(tokens, start_pos, last);
+    batch_all_out_ = saved;
+    return ok;
+}
+
 bool Q4nxNpuForward::step_batch(const std::vector<int>& tokens, int start_pos,
                                 std::vector<float>& logits) {
     if (!ready_) { err_ = "not initialised"; return false; }
@@ -3097,11 +3107,23 @@ bool Q4nxNpuForward::step_batch(const std::vector<int>& tokens, int start_pos,
     // The next-token logits come from the last row of the batch, after the final
     // norm -- step() applies final_norm_ before run_lm_head, and feeding the raw
     // residual here produced garbage logits (the M=1 self-check caught it).
-    std::vector<float> last(X.begin() + (size_t)(M - 1) * H, X.end());
     std::vector<float> fnorm(H);
-    rmsnorm(last.data(), final_norm_.data(), H, fnorm.data());
+    bool ok = true;
+    if (batch_all_out_) {
+        // step_batch_all(): the verifier needs the model's own choice at every drafted position, so every
+        // row is normed and the head runs ONCE for the batch (the M=16 twin streams its weights once).
+        std::vector<float> norm_all((size_t)M * (size_t)H);
+        for (int m = 0; m < M; m++)
+            rmsnorm(&X[(size_t)m * H], final_norm_.data(), H, &norm_all[(size_t)m * H]);
+        std::copy(norm_all.begin() + (size_t)(M - 1) * H, norm_all.end(), fnorm.begin());
+        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());
+    } else {
+        std::vector<float> last(X.begin() + (size_t)(M - 1) * H, X.end());
+        rmsnorm(last.data(), final_norm_.data(), H, fnorm.data());
+        ok = run_lm_head(fnorm, logits);
+    }
     last_hidden_ = fnorm;  // for last_hidden(): the same vector the head reads, last position only
-    const bool ok = run_lm_head(fnorm, logits);
     if (getenv("NPU_INFER_DUMP_LAYERS")) {
         int best[5] = {-1, -1, -1, -1, -1};
         for (int i = 0; i < (int)logits.size(); i++) {
@@ -4193,6 +4215,89 @@ bool Q4nxNpuForward::step(int token, int pos, std::vector<float>& logits) {
 }
 
 // tiled LM head: logits[K*t + n] = dot(embed[K*t + n], h); K = H, tile = QKV N
+bool Q4nxNpuForward::run_lm_head_rows(const std::vector<float>& h_rows, int M,
+                                     std::vector<float>& logits_all) {
+    const int H = c_.H, NV = c_.NV;
+    if (M <= 0) { err_ = "lm head rows: empty batch"; return false; }
+    if (M == 1 || h_rows.size() != (size_t)M * (size_t)H) {
+        std::vector<float> one(h_rows.begin(), h_rows.begin() + H);
+        std::vector<float> lg;
+        if (!run_lm_head(one, lg)) return false;
+        logits_all = std::move(lg);
+        return true;
+    }
+    logits_all.assign((size_t)M * (size_t)NV, 0.0f);
+    // One single-row call first: it prepares the packed tiles / f32 gather and refuses loudly (missing
+    // images, two-plane tile on a one-plane ELF) exactly as the token path does, so this function never
+    // has to duplicate that setup.  Its logits are the last row's.
+    {
+        std::vector<float> last(h_rows.begin() + (size_t)(M - 1) * H, h_rows.begin() + (size_t)M * H);
+        std::vector<float> lg;
+        if (!run_lm_head(last, lg)) return false;
+        std::copy(lg.begin(), lg.end(), logits_all.begin() + (size_t)(M - 1) * (size_t)NV);
+    }
+    const bool rows_supported = !host_gemm_ && head_pb16i_ && ctx_head_ != nullptr && !lm_head_packed_.empty();
+    if (!rows_supported) {
+        // host_gemm and the int8 head have no per-row launcher; the rows are rare (speculative verify only),
+        // so they go one at a time rather than growing a second code path in the hot token routine.
+        for (int m = 0; m + 1 < M; m++) {
+            std::vector<float> one(h_rows.begin() + (size_t)m * H, h_rows.begin() + (size_t)(m + 1) * H);
+            std::vector<float> lg;
+            if (!run_lm_head(one, lg)) return false;
+            std::copy(lg.begin(), lg.end(), logits_all.begin() + (size_t)m * (size_t)NV);
+        }
+        return true;
+    }
+    // The pb16i head, once for the whole batch: per-row rotation + per-row activation scale (the launcher and
+    // the dequantiser both carry per-row scales), one launch per tile, then the host finishing each row.
+    int i_qkv = -1;
+    for (int d = 0; d < ndesigns_; d++)
+        if (designs_[d].K == H && designs_[d].N > 0 && (i_qkv < 0 || designs_[d].N > designs_[i_qkv].N))
+            i_qkv = d;
+    if (i_qkv < 0) { err_ = "LM head rows: no GEMM design with K=H available"; return false; }
+    const int T = designs_[i_qkv].N;
+    const int ntiles = (NV + T - 1) / T;
+    if ((int)lm_head_packed_.size() != ntiles) { err_ = "LM head rows: packed tiles not resident"; return false; }
+    auto hsig = had_signs_.find(H);
+    if (hsig == had_signs_.end()) {
+        err_ = "pb16i head rows: no Hadamard sign vector for K=" + std::to_string(H);
+        return false;
+    }
+    std::vector<float> hr((size_t)M * (size_t)H), scales((size_t)M);
+    for (int m = 0; m < M; m++) {
+        const float* src = &h_rows[(size_t)m * H];
+        float* dst = &hr[(size_t)m * H];
+        std::copy(src, src + H, dst);
+        hadamard_rotate(hsig->second.data(), dst, H, had_block_);
+        float amax = 0.0f;
+        for (int k = 0; k < H; k++) {
+            const float a = std::fabs(dst[(size_t)k]);
+            if (a > amax) amax = a;
+        }
+        scales[(size_t)m] = amax > 0.0f ? amax / 127.0f : 1.0f;
+    }
+    I8Ctx& lctx = *ctx_head_;
+    std::vector<float> got((size_t)M * (size_t)T);
+    for (int t = 0; t < ntiles; t++) {
+        const int n0 = t * T, n1 = std::min(NV, n0 + T);
+        xrt::run run = lctx.launch_async_with_bo_rows(*lm_head_packed_[(size_t)t].bo, hr.data(), M, H,
+                                                      scales.data());
+        lctx.wait_kernel(run);
+        lctx.readback();
+        std::fill(got.begin(), got.end(), 0.0f);
+        lctx.dequant_pb16i(got.data(), M, T, scales.data(), lm_head_packed_[(size_t)t].eref,
+                           ctx_head_->pb16i_planes);
+        for (int m = 0; m < M; m++) {
+            for (int n = n0; n < n1; n++) {
+                float v = got[(size_t)m * T + (n - n0)];
+                if (sem_.logit_softcap > 0) v = sem_.logit_softcap * std::tanh(v / sem_.logit_softcap);
+                logits_all[(size_t)m * (size_t)NV + (size_t)n] = v;
+            }
+        }
+    }
+    return true;
+}
+
 bool Q4nxNpuForward::run_lm_head(const std::vector<float>& h, std::vector<float>& logits) {
     const int H = c_.H, NV = c_.NV;
     if (host_gemm_) {

tests/mtp_batch_all_test.cpp

@@ -0,0 +1,170 @@
+// 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.
+
+// mtp_batch_all_test.cpp -- step_batch_all() must return, for every row, the logits a sequential decode
+// would have produced at that position.  That is the whole contract speculative verification rests on: a
+// draft token is accepted only when the model's OWN choice at that position matches it, so a per-row
+// logit that is merely "close" is not enough -- the argmax has to agree.
+//
+// Host-only: the forward runs in host-GEMM mode (no device, no XRT).  Needs a Q4NX container, so it is
+// skipped (return 77) unless ONEBIT_TEST_Q4NX points at one:
+//
+//     ONEBIT_TEST_Q4NX=~/models/npu/Qwen3-0.6B ./mtp_batch_all_test
+//
+// Checked here:
+//   [1] row m of the batch equals the logits of a sequential step at position start+m (rel diff reported,
+//       argmax must match exactly);
+//   [2] the last row equals what plain step_batch() returns -- the batched prefill path is untouched;
+//   [3] the argmax of every row is stable across a repeat run (nothing depends on call order).
+#include "q4nx_forward.h"
+
+#include <cmath>
+#include <cstdio>
+#include <cstdlib>
+#include <string>
+#include <vector>
+
+// Q4nxNpuForward is at global scope in this header (npu/forward/include/q4nx_forward.h).
+
+namespace {
+
+double rel_diff(const std::vector<float>& a, const std::vector<float>& b, double* maxabs = nullptr) {
+    if (a.size() != b.size() || a.empty()) return 1e9;
+    double num = 0, den = 0, ma = 0;
+    for (size_t i = 0; i < a.size(); i++) {
+        const double d = std::fabs((double)a[i] - (double)b[i]);
+        num += d;
+        den += std::fabs((double)a[i]);
+        if (a[i] > ma) ma = a[i];
+        if (maxabs && d > *maxabs) *maxabs = d;
+    }
+    return num / (den + 1e-30);
+}
+
+int argmax(const std::vector<float>& v) {
+    int best = 0;
+    for (size_t i = 1; i < v.size(); i++)
+        if (v[i] > v[best]) best = (int)i;
+    return best;
+}
+
+int fails = 0;
+void check(bool cond, const std::string& what) {
+    std::printf("%s  %s\n", cond ? "PASS" : "FAIL", what.c_str());
+    if (!cond) fails++;
+}
+
+}  // namespace
+
+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 to run this\n");
+        return 77;
+    }
+    const std::string model = std::string(dir) + "/model.q4nx";
+    Q4nxNpuForward fwd;
+    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;   // a container this build cannot read is not a failure of step_batch_all
+    }
+    std::fprintf(stderr, "[test] phase 0: one token for V\n");
+    const int V = [&] {
+        std::vector<float> lg;
+        fwd.reset();
+        if (!fwd.step(1, 0, lg)) {
+            std::printf("SKIP: step() failed (%s)\n", fwd.last_error());
+            std::exit(77);
+        }
+        return (int)lg.size();
+    }();
+    std::printf("container %s, vocab %d\n", model.c_str(), V);
+
+    // A short batch of plausible tokens.  Values do not matter: the check is batched-vs-sequential.
+    // Two rows is enough to check the contract and keeps a host-GEMM run of a 27B container bounded.
+    const std::vector<int> batch = {1, 2 % V};
+
+    // Sequential reference: reset, then feed the same tokens one at a time.
+    std::fprintf(stderr, "[test] phase 1: sequential steps\n");
+    fwd.reset();
+    std::vector<std::vector<float>> seq;
+    for (size_t m = 0; m < batch.size(); m++) {
+        std::vector<float> lg;
+        if (!fwd.step(batch[m], (int)m, lg)) {
+            std::printf("FAIL  sequential step %zu failed: %s\n", m, fwd.last_error());
+            return 1;
+        }
+        seq.push_back(lg);
+    }
+
+    // [1] the batch, all rows
+    std::fprintf(stderr, "[test] phase 2: step_batch_all\n");
+    fwd.reset();
+    std::vector<float> all;
+    if (!fwd.step_batch_all(batch, 0, all)) {
+        std::printf("FAIL  step_batch_all failed: %s\n", fwd.last_error());
+        return 1;
+    }
+    if ((int)all.size() != (int)batch.size() * V) {
+        std::printf("FAIL  step_batch_all returned %zu floats, want %zu\n", all.size(), batch.size() * (size_t)V);
+        return 1;
+    }
+    for (size_t m = 0; m < batch.size(); m++) {
+        std::vector<float> row(all.begin() + m * (size_t)V, all.begin() + (m + 1) * (size_t)V);
+        double maxabs = 0;
+        const double rel = rel_diff(row, seq[m], &maxabs);
+        char msg[200];
+        std::snprintf(msg, sizeof msg,
+                      "row %zu: batched vs sequential rel-L1 %.2e (max|d| %.2e), argmax %d vs %d", m, rel,
+                      maxabs, argmax(row), argmax(seq[m]));
+        // the argmax is the contract; the values must be the same computation, so the tolerance is tight
+        check(argmax(row) == argmax(seq[m]) && rel < 5e-3, msg);
+    }
+
+    // [2] the last row is exactly what plain step_batch() returns
+    std::fprintf(stderr, "[test] phase 3: step_batch\n");
+    fwd.reset();
+    std::vector<float> tail;
+    if (!fwd.step_batch(batch, 0, tail)) {
+        std::printf("FAIL  step_batch failed: %s\n", fwd.last_error());
+        return 1;
+    }
+    {
+        std::vector<float> lastrow(all.begin() + (batch.size() - 1) * (size_t)V, all.end());
+        double maxabs = 0;
+        const double rel = rel_diff(lastrow, tail, &maxabs);
+        char msg[200];
+        std::snprintf(msg, sizeof msg, "last row == step_batch(): rel-L1 %.2e, argmax %d vs %d", rel,
+                      argmax(lastrow), argmax(tail));
+        check(argmax(lastrow) == argmax(tail) && rel < 1e-6, msg);
+    }
+
+    // [3] repeatability
+    std::fprintf(stderr, "[test] phase 4: repeat\n");
+    fwd.reset();
+    std::vector<float> again;
+    if (fwd.step_batch_all(batch, 0, again)) {
+        double maxabs = 0;
+        const double rel = rel_diff(all, again, &maxabs);
+        char msg[200];
+        std::snprintf(msg, sizeof msg, "repeat run identical: rel-L1 %.2e", rel);
+        check(rel == 0.0, msg);
+    } else {
+        check(false, "repeat step_batch_all failed");
+    }
+
+    std::printf(fails ? "\n%d CHECK(S) FAILED\n" : "\nall checks passed\n", fails);
+    return fails ? 1 : 0;
+}