diff --git a/CMakeLists.txt b/CMakeLists.txt index 46fc02c49..6dc453372 100644 --- a/CMakeLists.txt +++ b/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) diff --git a/npu/forward/include/spec_greedy.h b/npu/forward/include/spec_greedy.h new file mode 100644 index 000000000..ec8dc8c7b --- /dev/null +++ b/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 +#include + +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 recent_; + bool off_ = false; +}; + +} // namespace onebit::npu diff --git a/tests/spec_greedy_test.cpp b/tests/spec_greedy_test.cpp new file mode 100644 index 000000000..6fd9b2c1f --- /dev/null +++ b/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 + +#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; +} diff --git a/tools/dflash2_probe.cpp b/tools/dflash2_probe.cpp index 3abe35886..571258389 100644 --- a/tools/dflash2_probe.cpp +++ b/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 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,15 +532,14 @@ 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 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; @@ -552,7 +547,7 @@ int loop(int argc, char** argv) { 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