npu: the greedy speculative round and its fallback rule as a pure header (refs #1781)

CMakeLists.txt

@@ -1694,6 +1694,11 @@ 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)
+# The greedy speculative round and its fallback rule (npu/forward/include/spec_greedy.h, engine#1781). Header-only.
+add_executable(spec_greedy_test tests/spec_greedy_test.cpp)
+target_include_directories(spec_greedy_test PRIVATE npu/forward/include)
+add_test(NAME spec_greedy COMMAND spec_greedy_test)
+add_dependencies(onebit spec_greedy_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)

npu/forward/include/spec_greedy.h

@@ -0,0 +1,81 @@
+// 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.
+
+// spec_greedy.h -- the two decisions of a greedy speculative round, pure (engine#1781).
+//
+// A round drafts k tokens after the last emitted one, the target verifies [last, d1..dk] in one batched pass, and
+// row i of that pass is the target's own choice for the token after position i. So:
+//   - accepted = the longest prefix of the drafts that equals the rows' choices: d_{i+1} is kept while row i chose it;
+//   - next     = row `accepted`'s choice, the target's own token after the last kept one (always defined, because the
+//                pass has k + 1 rows).
+// The round emits accepted + 1 tokens, and the caller commits accepted + 1 rows of the verify pass.
+//
+// The fallback: a request whose last `window` rounds emitted fewer than `below` tokens on average stops speculating
+// and finishes with plain greedy steps. On the served 27B a round costs about 3.7 plain steps, so `below` = 3.7 is
+// break-even there (engine#1781). It looks only at the last rounds, so early luck does not keep a request speculating.
+
+#pragma once
+
+#include <cstddef>
+#include <vector>
+
+namespace onebit::npu {
+
+struct GreedyRound {
+    int accepted = 0;   // drafts kept, 0..k
+    int next = -1;      // the target's own token after them
+};
+
+// drafts[0..k), choices[0..k] (k + 1 rows: the target's argmax per verify row). Returns accepted = -1 when the pass
+// has fewer than k + 1 rows or k < 0.
+inline GreedyRound greedy_round(const int* drafts, int k, const int* choices, int rows) {
+    GreedyRound r;
+    if (k < 0 || rows < k + 1 || (k > 0 && !drafts) || !choices) {
+        r.accepted = -1;
+        return r;
+    }
+    int a = 0;
+    while (a < k && choices[a] == drafts[a]) a++;
+    r.accepted = a;
+    r.next = choices[a];
+    return r;
+}
+
+class SpecFallback {
+public:
+    // below <= 0 or window < 1: never falls back.
+    SpecFallback(double below, int window) : below_(below), window_(window) {}
+    // Record one round's emitted tokens (accepted + 1). Returns true once the request should stop speculating; stays
+    // true from then on.
+    bool record(int emitted) {
+        if (off_) return true;
+        recent_.push_back(emitted);
+        if (below_ <= 0.0 || window_ < 1 || int(recent_.size()) < window_) return false;
+        double sum = 0.0;
+        for (size_t i = recent_.size() - size_t(window_); i < recent_.size(); i++) sum += recent_[i];
+        off_ = sum / double(window_) < below_;
+        return off_;
+    }
+    bool off() const { return off_; }
+    int rounds() const { return int(recent_.size()); }
+
+private:
+    double below_;
+    int window_;
+    std::vector<int> recent_;
+    bool off_ = false;
+};
+
+}  // namespace onebit::npu

tests/spec_greedy_test.cpp

@@ -0,0 +1,79 @@
+// 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.
+
+// spec_greedy_test.cpp -- npu/forward/include/spec_greedy.h: the greedy round and the fallback rule (engine#1781).
+
+#include <cstdio>
+
+#include "spec_greedy.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 d[3] = {5, 6, 7};
+        const int all[4] = {5, 6, 7, 9};
+        GreedyRound r = greedy_round(d, 3, all, 4);
+        check(r.accepted == 3 && r.next == 9, "every draft kept: the 4th row gives the token after them");
+        const int mid[4] = {5, 4, 7, 9};
+        r = greedy_round(d, 3, mid, 4);
+        check(r.accepted == 1 && r.next == 4, "a miss in the middle: one kept, the row that missed gives the next token, "
+                                              "and a later match does not count");
+        const int none[4] = {8, 6, 7, 9};
+        r = greedy_round(d, 3, none, 4);
+        check(r.accepted == 0 && r.next == 8, "the first draft missed: nothing kept, row 0's choice is next");
+        r = greedy_round(nullptr, 0, all, 1);
+        check(r.accepted == 0 && r.next == 5, "no drafts: a plain step through the same rule");
+        check(greedy_round(d, 3, all, 3).accepted == -1, "a pass without the k + 1-th row is refused");
+        check(greedy_round(d, -1, all, 4).accepted == -1 && greedy_round(d, 3, nullptr, 4).accepted == -1,
+              "a negative k or no choices is refused");
+    }
+    {
+        SpecFallback f(3.7, 3);
+        check(!f.record(8) && !f.record(8), "fewer rounds than the window: never falls back");
+        check(!f.record(1), "8, 8, 1: average 5.67, above 3.7");
+    }
+    {
+        SpecFallback f(3.7, 3);
+        f.record(8); f.record(8); f.record(1);
+        check(f.record(2), "the window is the LAST three rounds: 8, 1, 2 averages 3.67, under 3.7");
+        f.record(8); f.record(8); f.record(8);
+        check(f.record(8) && f.off(), "once off, it stays off, even after rounds that would pass the window again");
+    }
+    {
+        SpecFallback f(3.7, 3);
+        bool off = false;
+        for (int i = 0; i < 10; i++) off = off || f.record(7);
+        check(!off && f.rounds() == 10, "a request that keeps accepting never falls back");
+        SpecFallback never(0.0, 3);
+        check(!never.record(1) && !never.record(1) && !never.record(1), "below <= 0 turns the rule off");
+    }
+    {
+        // The rule is "average BELOW the bar", strictly: a window that averages exactly the bar keeps speculating.
+        // 2, 2, 2 against 2.0 is exact in double arithmetic, so this pins the < against a <= without a rounding margin.
+        SpecFallback f(2.0, 3);
+        check(!f.record(2) && !f.record(2) && !f.record(2) && !f.off(), "2, 2, 2 against 2.0: exactly at the bar, stays on");
+        check(f.record(1), "then 2, 2, 1 (average 1.67): under the bar, falls back");
+    }
+    if (g_fail) { std::printf("spec_greedy: %d of %d checks FAILED\n", g_fail, g_n); return 1; }
+    std::printf("spec_greedy: ok (%d checks)\n", g_n);
+    return 0;
+}

tools/dflash2_probe.cpp

@@ -55,6 +55,7 @@
 #include "forward/include/q4nx_dequant.h"
 #include "pack_import.h"
 #include "q4nx_forward.h"
+#include "spec_greedy.h"     // the greedy round and the fallback rule
 #include "npu_engine_i8ctx_inc.h"   // I8Ctx: one int8 GEMM design from a full ELF
 
 namespace {
@@ -465,21 +466,16 @@ int loop(int argc, char** argv) {
     // rounds emitted fewer tokens than that on average stops speculating and finishes with plain greedy steps.
     const double fallback_below = getenv("DFLASH2_FALLBACK") ? std::atof(getenv("DFLASH2_FALLBACK")) : 0.0;
     const int fallback_after = getenv("DFLASH2_FALLBACK_ROUNDS") ? std::max(1, std::atoi(getenv("DFLASH2_FALLBACK_ROUNDS"))) : 3;
-    // The rule looks at the last DFLASH2_FALLBACK_ROUNDS rounds only, so a request whose first rounds happened to
-    // accept well does not keep speculating on that credit.
-    std::vector<int> recent;   // tokens emitted by each round, newest last
+    // The rule (spec_greedy.h) looks at the last DFLASH2_FALLBACK_ROUNDS rounds only, so a request whose first rounds
+    // happened to accept well does not keep speculating on that credit.
+    SpecFallback fallback(fallback_below, fallback_after);
     bool plain = false;
     long plain_steps = 0;
     while (int(out.size()) < want_tokens) {
-        if (!plain && fallback_below > 0.0 && int(recent.size()) >= fallback_after) {
-            double sum = 0.0;
-            for (size_t i = recent.size() - size_t(fallback_after); i < recent.size(); i++) sum += recent[i];
-            const double avg = sum / double(fallback_after);
-            if (avg < fallback_below) {
-                plain = true;
-                std::printf("loop: falling back to plain decoding after %ld rounds (last %d rounds: %.2f tokens per round, below %.2f)\n",
-                            rounds, fallback_after, avg, fallback_below);
-            }
+        if (!plain && fallback.off()) {
+            plain = true;
+            std::printf("loop: falling back to plain decoding after %ld rounds (last %d rounds below %.2f tokens per round)\n",
+                        rounds, fallback_after, fallback_below);
         }
         if (plain) {
             if (!fwd.step(cur, pos, logits)) die(std::string("step: ") + fwd.last_error());
@@ -536,23 +532,22 @@ int loop(int argc, char** argv) {
         fwd.begin_verify(int(vt.size()), pos);
         if (!fwd.step_batch_all(vt, pos, logits_all)) die(std::string("verify: ") + fwd.last_error());
         const size_t NV = logits_all.size() / vt.size();
-        int a = 0;
-        int next = -1;
-        for (int i = 0; i < int(vt.size()); i++) {
-            const float* lr = logits_all.data() + size_t(i) * NV;
-            const int g = int(std::max_element(lr, lr + NV) - lr);
-            if (i < per && g == drafts[size_t(i)]) { a++; continue; }
-            next = g;
-            break;
+        std::vector<int> choices(vt.size());
+        for (size_t i = 0; i < vt.size(); i++) {
+            const float* lr = logits_all.data() + i * NV;
+            choices[i] = int(std::max_element(lr, lr + NV) - lr);
         }
+        const GreedyRound gr = greedy_round(drafts.data(), per, choices.data(), int(choices.size()));
+        if (gr.accepted < 0) die("verify: the pass returned fewer rows than it was given");
+        const int a = gr.accepted, next = gr.next;
         if (!fwd.commit_verify_rows(a + 1)) die(std::string("commit: ") + fwd.last_error());
         taps.resize(tap_base + size_t(a + 1) * rowf);   // keep the rows of the committed positions only
         t_verify += now_s() - c0;
         // emit
         for (int i = 0; i < a; i++) out.push_back(drafts[size_t(i)]);
         out.push_back(next);
         rounds++;
-        recent.push_back(a + 1);
+        fallback.record(a + 1);
         accepted_total += a;
         hist[size_t(a)]++;
         // extend the fused context by the committed rows