npu: the speculative round's seam and bookkeeping, with a rejection in the middle pinned (refs #1781; supersedes #2059)

CMakeLists.txt

@@ -507,6 +507,18 @@ add_test(NAME pack_initializer_shape COMMAND pack_initializer_shape_test)
 add_dependencies(onebit pack_initializer_shape_test)
 target_include_directories(npu_sampling_test PRIVATE app)
 target_link_libraries(npu_sampling_test PRIVATE nlohmann_json::nlohmann_json)
+# The speculative round's bookkeeping (app/spec_decode.h, engine#1781): a drafter's proposals against the verifier's
+# per-row logits, with the acceptance rule from npu/mtp/spec_accept.h (#1958, merged). Header-only and pure -- the
+# uniforms come from the caller, so every arm is a hand-computed value and there is no model, device or RNG here.
+# `add_dependencies(onebit ...)` because a target ci.yml never builds is the failure that check exists to catch.
+add_executable(spec_decode_test tests/spec_decode_test.cpp)
+# ${CMAKE_SOURCE_DIR} as well as app/, because app/spec_decode.h includes "npu/mtp/spec_accept.h" by its repo-relative
+# path -- which is how every other consumer of that header spells it. Without the root on this target's path the build
+# fails with "npu/mtp/spec_accept.h: No such file or directory", measured on CI run 42935896650 (the risk the PR body
+# named, so it cost one run rather than a review).
+target_include_directories(spec_decode_test PRIVATE app ${CMAKE_SOURCE_DIR})
+add_test(NAME spec_decode COMMAND spec_decode_test)
+add_dependencies(onebit spec_decode_test)
 # The reader for a tools/moe_placement.py table (app/moe_placement.h, engine#1702): a consumer that splits experts across
 # devices adopts a table only if it was written for this model -- schema, group count, every routed layer, every expert
 # id exactly once -- and only on the held-out share, never the in-sample one. Header only, no device, no model.

app/spec_decode.h

@@ -0,0 +1,99 @@
+// 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_decode.h -- the speculative round: what a drafter has to hand the verifier, and the decision that turns
+// M x NV target logits plus a draft's own distribution into the tokens a round emits (engine#1781, the wiring task).
+//
+// The pieces this does NOT own, deliberately, because they live elsewhere and are already measured:
+//   * the acceptance rule itself -- npu/mtp/spec_accept.h (#1958, merged), which reduces to "accept iff the draft is
+//     the target's argmax" at temperature 0;
+//   * the verify pass -- app/forward_serve.cpp's begin_verify()/step_batch_all(), which already returns one row of
+//     logits per position (`M x NV`), so `p` needs no new plumbing;
+//   * the drafter -- #1780's DFlash packages and #1957's gate, which will implement DraftProvider.
+//
+// What it does own is the bookkeeping that a speedup number rests on: which drafts were accepted, in what order, where
+// the chain stopped, and how many were proposed -- because a tok/s figure without the acceptance that produced it is
+// not reproducible. The decision function is pure and takes its uniforms from the caller, so the round can be tested
+// with no model, no device and no RNG, and so the two draws actually come from the sampler's stream rather than from a
+// per-position generator (which would follow `p` in distribution but not in *stream*).
+#pragma once
+
+#include <cstddef>
+#include <vector>
+
+#include "npu/mtp/spec_accept.h"
+
+namespace onebit::spec {
+
+// One drafter answer: up to k proposed ids, and the drafter's own distribution at each of them, row-major k x V.
+struct DraftStep {
+    std::vector<int> ids;
+    std::vector<float> q;
+    int rows() const { return int(ids.size()); }
+    bool valid(int V) const { return V > 0 && q.size() == ids.size() * std::size_t(V); }
+};
+
+// What a speculative loop needs from a drafter, and nothing else. Returning false means "nothing to say" (no
+// drafter, out of range, or a refusal), and the caller then takes exactly the non-speculative path it took before.
+class DraftProvider {
+  public:
+    virtual ~DraftProvider() = default;
+    virtual bool propose(const std::vector<int>& ctx, int k, int V, DraftStep& out) = 0;
+};
+
+// One round's outcome. `emitted` is the accepted drafts in order, plus the resampled token when the chain stops on a
+// rejection -- never more than drafts.size() + 0: the resample *replaces* the rejected draft, so a round of k drafts
+// emits between 1 and k tokens. When every draft is accepted the round emits exactly k and no more: the token after
+// them comes from the verifier's (k+1)-th row, which the CALLER samples (spec_accept_prefix's contract). A loop built
+// on this must not drop that sample -- it is one token per fully-accepted round.
+struct SpecRound {
+    std::vector<int> emitted;
+    int accepted = 0;
+    int proposed = 0;
+    bool stopped_on_rejection = false;
+};
+
+// Decide one round. `p_rows` and `q_rows` are drafts.size() x V, row-major, both aligned with `drafts`. `ua`/`ur` are
+// drafts.size() uniforms in [0,1) for the two draws each decision needs, taken from the caller's stream in order: if
+// either is null, or either is outside [0,1), the round refuses (returns proposed == 0) rather than drawing for itself.
+// A draft whose q(x) is 0 rejects and resamples (#1958's premise: `x` was not drawn from q), which is the case that
+// keeps the division out of the rule.
+inline SpecRound decide(const float* p_rows, const float* q_rows, int V, const std::vector<int>& drafts,
+                        const double* ua, const double* ur) {
+    SpecRound out;
+    if (!p_rows || !q_rows || V <= 0 || drafts.empty() || !ua || !ur) return out;
+    const std::size_t n = drafts.size();
+    for (std::size_t d = 0; d < n; ++d) {
+        const double a = ua[d], r = ur[d];
+        if (!(a >= 0.0 && a < 1.0) || !(r >= 0.0 && r < 1.0)) return SpecRound{};  // refuse, do not draw here
+    }
+    out.proposed = int(n);
+    const float* p = p_rows;
+    const float* q = q_rows;
+    for (std::size_t d = 0; d < n; ++d) {
+        const npu::mtp::SpecStep s =
+            npu::mtp::spec_accept_step(p + d * std::size_t(V), q + d * std::size_t(V), V, drafts[d], ua[d], ur[d]);
+        out.emitted.push_back(s.token);
+        if (s.accepted) {
+            ++out.accepted;
+            continue;
+        }
+        out.stopped_on_rejection = true;
+        break;  // the chain stops at the first rejection: every later draft was conditioned on a token that was not chosen
+    }
+    return out;
+}
+
+}  // namespace onebit::spec

tests/spec_decode_test.cpp

@@ -0,0 +1,121 @@
+// 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_decode_test -- the speculative round's bookkeeping (engine#1781). No model, no device, no RNG: `decide()` is
+// pure and takes its uniforms from the caller, so every arm below is a hand-computed expected value rather than a
+// tolerance. Five arms, and (v) is the one that matters most: with no provider the loop must behave exactly as it did
+// before the wiring existed.
+//
+//   (i)   a partially-accepted round: emitted = accepted ++ resample, counts exact
+//   (ii)  the chain stops at the FIRST rejection, so later drafts are never emitted
+//         (ii-b puts the rejection in the middle of three drafts, where a later draft exists to be wrongly emitted)
+//   (iii) a fully-accepted round emits all k and resamples nothing
+//   (iv)  q(x) == 0 rejects and resamples (the premise `x was drawn from q` is false) instead of dividing
+//   (v)   the no-drafter path is the caller's untouched branch -- asserted here as "decide() refuses with no drafts",
+//         which is the condition the loop keys on
+#include "spec_decode.h"
+
+#include <cmath>
+#include <cstdio>
+#include <vector>
+
+static int failures = 0;
+static void check(bool ok, const char* what) {
+    std::printf("%s  %s\n", ok ? "ok  " : "FAIL", what);
+    if (!ok) ++failures;
+}
+
+// Two-row fixtures at V = 4. p is the target, q the drafter. Both rows are distributions that sum to 1.
+static const float kP[2][4] = {{0.70f, 0.10f, 0.10f, 0.10f}, {0.10f, 0.70f, 0.10f, 0.10f}};
+static const float kQ[2][4] = {{0.10f, 0.70f, 0.10f, 0.10f}, {0.70f, 0.10f, 0.10f, 0.10f}};
+
+int main() {
+    const int V = 4;
+
+    // (i) two drafts, both accepted: p[x]=0.70 against q[x]=0.10, so min(1,p/q) is 1 and u<1 always accepts.
+    {
+        const std::vector<int> drafts{0, 1};
+        const double ua[2] = {0.99, 0.99}, ur[2] = {0.5, 0.5};
+        const onebit::spec::SpecRound r = onebit::spec::decide(&kP[0][0], &kQ[0][0], V, drafts, ua, ur);
+        check(r.proposed == 2 && r.accepted == 2 && !r.stopped_on_rejection && r.emitted.size() == 2 &&
+                  r.emitted[0] == 0 && r.emitted[1] == 1,
+              "(i) both accepted -> emitted = the two drafts, nothing resampled");
+    }
+    // (ii) first draft accepted, second rejected. Both rows propose token 0: row 0 has p=0.70,q=0.10 so min(1,p/q)=1
+    // (accepts for any u), row 1 has p=0.10,q=0.70 so min(1,p/q)=1/7 and u=0.99 > 1/7 rejects. Proposing token 1 on row 1
+    // would have been accepted (p/q = 7) -- which is what the first version of this fixture did, and the arm said so.
+    {
+        const std::vector<int> drafts{0, 0};
+        const double ua[2] = {0.01, 0.99}, ur[2] = {0.5, 0.5};
+        const onebit::spec::SpecRound r = onebit::spec::decide(&kP[0][0], &kQ[0][0], V, drafts, ua, ur);
+        check(r.proposed == 2 && r.accepted == 1 && r.stopped_on_rejection && r.emitted.size() == 2,
+              "(ii) one accepted then a rejection -> emitted = accepted ++ the resample");
+        check(r.emitted[0] == 0 && r.emitted[1] >= 0 && r.emitted[1] < V,
+              "(ii) emitted[0] is the accepted draft and emitted[1] is a valid resampled token");
+    }
+    // (ii-b) the rejection in the MIDDLE of three drafts. Arm (ii) rejects at its last draft, so nothing came after the
+    // rejection for a missing `break` to emit, and a decide() that kept going passed it (review of 030070e). Here the
+    // third draft would be accepted on its own row (p/q = 7); it must not be emitted, because it was conditioned on
+    // the second draft, which was not chosen.
+    {
+        const float P[3][4] = {{0.70f, 0.10f, 0.10f, 0.10f}, {0.10f, 0.70f, 0.10f, 0.10f}, {0.70f, 0.10f, 0.10f, 0.10f}};
+        const float Q[3][4] = {{0.10f, 0.70f, 0.10f, 0.10f}, {0.70f, 0.10f, 0.10f, 0.10f}, {0.10f, 0.70f, 0.10f, 0.10f}};
+        const std::vector<int> drafts{0, 0, 0};
+        const double ua[3] = {0.01, 0.99, 0.01}, ur[3] = {0.5, 0.5, 0.5};
+        const onebit::spec::SpecRound r = onebit::spec::decide(&P[0][0], &Q[0][0], V, drafts, ua, ur);
+        check(r.proposed == 3 && r.accepted == 1 && r.stopped_on_rejection && r.emitted.size() == 2,
+              "(ii-b) a rejection in the middle ends the round: one accepted, the resample, and nothing after it");
+    }
+    // (iii) three drafts at a row where p == q: min(1, p/q) == 1, so every u accepts and nothing is resampled.
+    {
+        const float same[3][4] = {{0.25f, 0.25f, 0.25f, 0.25f}, {0.25f, 0.25f, 0.25f, 0.25f},
+                                  {0.25f, 0.25f, 0.25f, 0.25f}};
+        const std::vector<int> drafts{0, 1, 2};
+        const double ua[3] = {0.9, 0.9, 0.9}, ur[3] = {0.5, 0.5, 0.5};
+        const onebit::spec::SpecRound r = onebit::spec::decide(&same[0][0], &same[0][0], V, drafts, ua, ur);
+        check(r.proposed == 3 && r.accepted == 3 && r.emitted.size() == 3 && !r.stopped_on_rejection,
+              "(iii) p == q -> all k accepted, no resample, emitted.size() == k");
+    }
+    // (iv) q(x) == 0: the draft was not drawn from q, so the step must reject and resample rather than divide by zero.
+    {
+        const float q0[1][4] = {{0.0f, 0.5f, 0.5f, 0.0f}};
+        const float p0[1][4] = {{0.9f, 0.05f, 0.05f, 0.0f}};
+        const std::vector<int> drafts{0};
+        const double ua[1] = {0.0}, ur[1] = {0.0};
+        const onebit::spec::SpecRound r = onebit::spec::decide(&p0[0][0], &q0[0][0], V, drafts, ua, ur);
+        check(r.proposed == 1 && r.accepted == 0 && r.stopped_on_rejection && r.emitted.size() == 1,
+              "(iv) q(x) == 0 rejects instead of dividing, and still emits one resampled token");
+    }
+    // (v) the no-drafter path: no drafts is exactly the condition the loop keys the speculative branch on, so decide()
+    //     must refuse for it. The caller's argmax branch is untouched by this header -- that is the guarantee.
+    {
+        const double ua[1] = {0.5}, ur[1] = {0.5};
+        const onebit::spec::SpecRound r = onebit::spec::decide(&kP[0][0], &kQ[0][0], V, std::vector<int>{}, ua, ur);
+        check(r.proposed == 0 && r.accepted == 0 && r.emitted.empty(),
+              "(v) no drafts -> refuses, so the loop keeps its own argmax path");
+    }
+    // and the seam's other refusal: uniforms outside [0,1) are refused, not clamped -- a caller that forgot to seed its
+    // RNG must not get a silent pass.
+    {
+        const std::vector<int> drafts{0};
+        const double bad[1] = {1.0}, ok[1] = {0.5};
+        const onebit::spec::SpecRound r1 = onebit::spec::decide(&kP[0][0], &kQ[0][0], V, drafts, bad, ok);
+        const onebit::spec::SpecRound r2 = onebit::spec::decide(&kP[0][0], &kQ[0][0], V, drafts, ok, bad);
+        check(r1.proposed == 0 && r2.proposed == 0, "a uniform of exactly 1.0 (or outside [0,1)) is refused, not clamped");
+    }
+
+    if (failures == 0) std::printf("spec_decode: ok\n");
+    return failures ? 1 : 0;
+}