diff --git a/CMakeLists.txt b/CMakeLists.txt index ce06e4a37..91daf0836 100644 --- a/CMakeLists.txt +++ b/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. diff --git a/app/spec_decode.h b/app/spec_decode.h new file mode 100644 index 000000000..a7f5d3cc5 --- /dev/null +++ b/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 +#include + +#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 ids; + std::vector 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& 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 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& 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 diff --git a/tests/spec_decode_test.cpp b/tests/spec_decode_test.cpp new file mode 100644 index 000000000..38739f6f7 --- /dev/null +++ b/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 +#include +#include + +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 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 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 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 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 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{}, 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 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; +}