From ee59a4327d147f579972198c640041c948312f5c Mon Sep 17 00:00:00 2001 From: bong-water-water-bong Date: Thu, 12 Nov 2026 01:34:32 -0400 Subject: [PATCH] npu: the DFlash draft as mtp::run()'s drafter (DFlashHead, refs #1781) npu/dflash_head.h: a DraftHead over dflash_forward with needs_last_hidden() == false, so mtp::run() takes the one-pass round. It reads the forward's tap sink, fuses committed rows into its DraftKvCache at draft time, and after a verify batch trims the sink to the rows the round committed (the forward writes a tap row for every position, rejected drafts included). Embedding and pick rule (argmax or a DFlash 2 selector) are hooks; a draft at a position the sink does not end at is refused with a reason. tests/dflash_head_test.cpp (ctest dflash_head): three pick rules x k = 2, 7 on a toy base that writes taps the same way: text equals plain decode, the sink holds exactly the committed rows, every draft equals draft_forward() over the committed rows, bit for bit. Six mutants each fail it. Co-Authored-By: Claude Opus 5.5 --- CMakeLists.txt | 8 ++ docs/npu.md | 8 ++ npu/dflash_head.h | 149 +++++++++++++++++++++++ tests/dflash_head_test.cpp | 241 +++++++++++++++++++++++++++++++++++++ 4 files changed, 406 insertions(+) create mode 100644 npu/dflash_head.h create mode 100644 tests/dflash_head_test.cpp diff --git a/CMakeLists.txt b/CMakeLists.txt index 73db2e28d..a4753c158 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -1449,6 +1449,14 @@ target_include_directories(dflash_forward_test PRIVATE npu npu/forward/include) target_link_libraries(dflash_forward_test PRIVATE onebit_npu) add_dependencies(onebit dflash_forward_test) add_test(NAME dflash_forward COMMAND dflash_forward_test ${ONEBIT_DFLASH_DRAFT}) +# engine#1781: the DFlash draft as mtp::run()'s drafter (npu/dflash_head.h) on a toy base that writes a tap row per +# position, rejected rows included: text equals plain decode, the sink keeps only committed rows, every draft equals +# draft_forward() over them. Host only. +add_executable(dflash_head_test tests/dflash_head_test.cpp npu/dflash_forward.cpp npu/dflash_draft.cpp npu/forward/src/q4nx_dequant.c) +target_include_directories(dflash_head_test PRIVATE npu npu/forward/include) +target_link_libraries(dflash_head_test PRIVATE onebit_npu) +add_dependencies(onebit dflash_head_test) +add_test(NAME dflash_head COMMAND dflash_head_test) add_executable(flm_head_import tools/flm_head_import.cpp) target_link_libraries(flm_head_import PRIVATE onebit_npu) add_executable(npu_flm_head_import_test tests/npu_flm_head_import_test.cpp) diff --git a/docs/npu.md b/docs/npu.md index 95187da73..1b026ccde 100644 --- a/docs/npu.md +++ b/docs/npu.md @@ -3374,6 +3374,14 @@ takes one pass per `k + 1` tokens and no step (k = 7: 0.125 base passes per toke and every committed row reaches the drafter once, in order. This is what the "one pass (verify only)" rows of [docs/oga-drafter-path.md](oga-drafter-path.md) price. +The DFlash draft is that head: `npu/dflash_head.h` (`DFlashHead`) drafts from the forward's tap sink +(`set_layer_taps`), fused into a `DraftKvCache` as rows are committed, and trims the sink after each verify batch to the +rows the round committed, because the forward writes a tap row for every position a pass runs, rejected drafts +included. The target's embedding and the rule that turns a drafted row into a token (the head's argmax or a DFlash 2 +selector) are hooks. `tests/dflash_head_test.cpp` runs it through `mtp::run()` on a toy base that writes taps the same +way: the text equals plain decode's, the sink ends holding exactly the committed rows, and every draft equals +`draft_forward()` over those rows, bit for bit. + **What would change the answer**: a cheaper verify batch (the array wait is ~60 % of a batched pass — a kernel lever) or a head that reliably produces a chain above 3.9 (widening `k` is not it: `k = 6` bought +0.05, and depth 6 measured 0 %). diff --git a/npu/dflash_head.h b/npu/dflash_head.h new file mode 100644 index 000000000..f835c862d --- /dev/null +++ b/npu/dflash_head.h @@ -0,0 +1,149 @@ +// 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. + +// dflash_head.h -- a DFlash draft as the speculative loop's drafter (engine#1781): mtp::DraftHead over dflash_forward. +// +// The draft reads the target's hidden states at its tap layers for every committed position (the forward writes +// them to a sink, Q4nxNpuForward::set_layer_taps), keeps them fused in a DraftKvCache, and drafts a block of +// [last token, mask, mask, ...]. It never reads the hidden after the last token, so needs_last_hidden() is false and +// mtp::run() verifies [t1, d1..dk] as one batch (docs/npu.md, "The draft loop, and whether to serve it"). +// +// The sink is the one piece of state shared with the forward, and the forward writes a row for EVERY position a pass +// runs, the rejected rows of a verify batch included. So: +// - at draft(), every row in the sink is a committed position: rows not yet in the cache are fused there; +// - commit_rows(M) after a verify batch keeps the M committed rows of that batch and drops the rest of its rows, +// which is what keeps a rejected draft's hidden states out of every later draft; +// - after a plain step (no draft that round) the one row the step wrote is committed, and nothing is dropped. +// The target's embedding table and the rule that turns a drafted row into a token (the target's head and argmax, or +// a DFlash 2 candidate selector) come in as hooks, so this is testable with no model and no device. +#pragma once + +#include +#include +#include +#include +#include +#include + +#include "dflash_forward.h" +#include "mtp/draft.h" + +namespace onebit::npu::dflash { + +struct HeadHooks { + // The target's embedding of `token`, hidden wide. False = no such token. + std::function& row)> embed; + // The draft token for one drafted row: `h` is the draft's output row (hidden wide), `prev` the token before it + // (the last emitted token for the first row). A negative return ends the chain there. + std::function pick; +}; + +class DFlashHead : public mtp::DraftHead { +public: + // `taps` is the forward's tap sink, one row of taps * hidden floats per position, starting at position 0. + // `block` is the draft's block size (so up to block - 1 drafts a round); `window` its sliding window, 0 = none. + DFlashHead(const DraftWeights& w, std::vector& taps, HeadHooks hooks, int mask_token, int block, + int window = 0) + : w_(w), taps_(taps), hooks_(std::move(hooks)), mask_token_(mask_token), block_(block), window_(window) {} + + bool needs_last_hidden() const override { return false; } + + int draft(const std::vector& /*hidden*/, int last_token, int pos, int k, std::vector& out, + std::vector* conf = nullptr) override { + out.clear(); + if (conf) conf->clear(); // no per-draft probability: mtp::run then skips its tau gate + drafted_ = false; + error_.clear(); + const size_t rowf = size_t(w_.taps) * size_t(w_.hidden); + if (rowf == 0 || taps_.size() % rowf != 0) return fail("the tap sink is not whole rows of taps * hidden"); + const int rows = int(taps_.size() / rowf); + // last_token runs at pos - 1, right after the committed rows: anything else means the sink and the loop + // disagree about the stream, and a draft from it would be conditioned on the wrong context. + if (pos != rows + 1) + return fail("drafting at position " + std::to_string(pos) + " over " + std::to_string(rows) + + " committed rows"); + const int n = std::min(k, block_ - 1); + if (n <= 0) return 0; + const size_t H = size_t(w_.hidden); + std::vector noise(size_t(block_) * H), row; + try { + if (cache_.rows < rows) { + const std::vector f = + fuse_context(w_, taps_.data() + size_t(cache_.rows) * rowf, rows - cache_.rows); + cache_extend(w_, cache_, f.data(), rows - cache_.rows); + } + if (mask_.empty() && (!hooks_.embed(mask_token_, mask_) || mask_.size() != H)) + return fail("no embedding for the mask token"); + if (!hooks_.embed(last_token, row) || row.size() != H) return fail("no embedding for the last token"); + std::copy(row.begin(), row.end(), noise.begin()); + for (int i = 1; i < block_; i++) std::copy(mask_.begin(), mask_.end(), noise.begin() + ptrdiff_t(size_t(i) * H)); + last_rows_ = draft_forward_cached(w_, cache_, noise.data(), block_, window_); + } catch (const std::exception& e) { + return fail(e.what()); + } + int prev = last_token; + for (int i = 1; i <= n; i++) { + const int t = hooks_.pick(last_rows_.data() + size_t(i) * H, prev); + if (t < 0) break; + out.push_back(t); + prev = t; + } + mark_ = taps_.size(); + drafted_ = true; + return int(out.size()); + } + + void commit_rows(const std::vector& /*tokens*/, const std::vector& /*hidden_rows*/, int M) override { + if (!drafted_) return; // a plain step: its one row is committed as written + drafted_ = false; + const size_t keep = mark_ + size_t(M) * size_t(w_.taps) * size_t(w_.hidden); + if (M < 0 || taps_.size() < keep) { + fail("commit of " + std::to_string(M) + " rows, but the pass wrote fewer"); + return; + } + taps_.resize(keep); + } + + void reset() override { + cache_ = DraftKvCache{}; + drafted_ = false; + error_.clear(); + } + + const DraftKvCache& cache() const { return cache_; } + // The draft's output rows of the last draft() ([block][hidden]) and the context rows it drafted over. + const std::vector& last_draft_rows() const { return last_rows_; } + int last_context_rows() const { return cache_.rows; } + // Why the last draft() returned nothing, or why the commit after it was refused; empty when neither went wrong. + const std::string& error() const { return error_; } + +private: + int fail(const std::string& why) { + error_ = why; + return 0; + } + + const DraftWeights& w_; + std::vector& taps_; + HeadHooks hooks_; + int mask_token_, block_, window_; + DraftKvCache cache_; + std::vector mask_, last_rows_; + size_t mark_ = 0; + bool drafted_ = false; + std::string error_; +}; + +} // namespace onebit::npu::dflash diff --git a/tests/dflash_head_test.cpp b/tests/dflash_head_test.cpp new file mode 100644 index 000000000..4dd572f1e --- /dev/null +++ b/tests/dflash_head_test.cpp @@ -0,0 +1,241 @@ +// 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. + +// dflash_head_test.cpp -- the DFlash draft as mtp::run()'s drafter (npu/dflash_head.h, engine#1781). Host only. +// +// The base is a toy that writes one tap row per position it runs -- the rejected rows of a verify batch too, as the +// forward's tap sink does -- and whose greedy choice is a fixed function of the token. The draft is a small random +// DraftWeights; its pick rule is the base's own rule (every draft right), that rule with every third draft wrong, or +// always wrong. Checked, for each, at k = 2 and 7 with a block of 8: +// [1] the text equals plain decode's; +// [2] the sink ends holding exactly the committed stream's rows (prompt plus every emitted token), bit for bit: no +// rejected row survives a commit; +// [3] every draft equals draft_forward() over the committed stream's rows alone, computed from scratch, bit for bit; +// [4] the right-everywhere rule has every draft accepted, the always-wrong one none; +// and that a draft at a position the sink does not end at is refused, with a reason. + +#include "dflash_head.h" + +#include +#include +#include +#include + +using onebit::npu::dflash::DFlashHead; +using onebit::npu::dflash::DraftLayerWeights; +using onebit::npu::dflash::DraftWeights; +using onebit::npu::dflash::HeadHooks; +using onebit::npu::mtp::BaseModel; +using onebit::npu::mtp::DraftHead; +using onebit::npu::mtp::Options; +using onebit::npu::mtp::Result; + +namespace { + +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++; +} + +constexpr int V = 16, H = 8, T = 2, BLOCK = 8; + +int next_of(int token) { return (token * 5 + 3) % V; } + +float hashf(uint64_t a, uint64_t b) { + uint64_t u = a * 0x9E3779B97F4A7C15ULL ^ (b + 0x632BE59BD9B4E019ULL); + u ^= u >> 29; u *= 0xBF58476D1CE4E5B9ULL; u ^= u >> 32; + return float(double(u >> 11) / double(1ULL << 53)) * 2.0f - 1.0f; +} +// The tap row the base writes for `token` run at `pos`. +void tap_row(int token, int pos, float* out) { + for (int i = 0; i < T * H; i++) out[i] = hashf(uint64_t(token) * 1000003ULL + uint64_t(pos), uint64_t(i)); +} + +class TapBase : public BaseModel { +public: + std::vector taps; // the sink + int vocab() const override { return V; } + int hidden_size() const override { return 1; } + int max_batch_rows() const override { return 16; } + bool step(int token, int pos, std::vector& logits) override { + if (pos != consumed_) return false; + put(token, pos); + consumed_++; + logits.assign(V, 0.0f); + logits[size_t(next_of(token))] = 1.0f; + hidden_.assign(1, float(consumed_)); + return true; + } + bool step_batch_all(const std::vector& tokens, int start, std::vector& all) override { + if (start != consumed_) return false; + all.assign(tokens.size() * size_t(V), 0.0f); + rows_.clear(); + for (size_t m = 0; m < tokens.size(); m++) { + put(tokens[m], start + int(m)); + all[m * size_t(V) + size_t(next_of(tokens[m]))] = 1.0f; + rows_.push_back(float(start + int(m) + 1)); + } + pending_ = int(tokens.size()); + consumed_ += pending_; + return true; + } + bool commit_verify(int a) override { // positions only; the sink is the drafter's to trim, as on the forward + if (a < 0 || a > pending_) return false; + consumed_ -= pending_ - a; + rows_.resize(size_t(a)); + pending_ = 0; + return true; + } + const std::vector& last_hidden() const override { return hidden_; } + const std::vector& last_hidden_rows() const override { return rows_; } + +private: + void put(int token, int pos) { + const size_t o = taps.size(); + taps.resize(o + size_t(T * H)); + tap_row(token, pos, taps.data() + o); + } + int consumed_ = 0, pending_ = 0; + std::vector hidden_, rows_; +}; + +DraftWeights random_model() { + DraftWeights w; + w.hidden = H; w.intermediate = 12; w.heads = 4; w.kv_heads = 2; w.head_dim = 4; w.taps = T; + w.rms_eps = 1e-6f; w.rope_theta = 10000.0f; + uint64_t c = 0; + auto vec = [&](size_t n, float s) { + std::vector v(n); + for (float& f : v) f = s * hashf(7, c++); + return v; + }; + auto gain = [&](size_t n) { + std::vector g = vec(n, 0.3f); + for (float& f : g) f += 1.0f; + return g; + }; + w.fc = vec(size_t(H) * H * T, 0.5f); + w.hidden_norm = gain(H); + w.norm = gain(H); + for (int l = 0; l < 2; l++) { + DraftLayerWeights L; + L.q = vec(16 * H, 0.6f); L.k = vec(8 * H, 0.6f); L.v = vec(8 * H, 0.6f); L.o = vec(H * 16, 0.4f); + L.gate = vec(12 * H, 0.5f); L.up = vec(12 * H, 0.5f); L.down = vec(H * 12, 0.4f); + L.input_norm = gain(H); L.post_norm = gain(H); L.q_norm = gain(4); L.k_norm = gain(4); + w.layers.push_back(L); + } + return w; +} + +void embed(int token, std::vector& row) { + row.resize(H); + for (int i = 0; i < H; i++) row[size_t(i)] = hashf(uint64_t(token) + 99, uint64_t(i)); +} + +// Records each draft's rows and the context it drafted over, so [3] can recompute them afterwards. +struct Logged : DraftHead { + DFlashHead& h; + std::vector ctx, last; + std::vector> out; + explicit Logged(DFlashHead& x) : h(x) {} + bool needs_last_hidden() const override { return h.needs_last_hidden(); } + int draft(const std::vector& hid, int lt, int pos, int k, std::vector& o, + std::vector* conf) override { + const int n = h.draft(hid, lt, pos, k, o, conf); + // The context a draft at `pos` must see is the pos - 1 committed rows before it, whatever the head reports. + if (n > 0) { ctx.push_back(pos - 1); last.push_back(lt); out.push_back(h.last_draft_rows()); } + return n; + } + void commit_rows(const std::vector& t, const std::vector& r, int M) override { h.commit_rows(t, r, M); } +}; + +} // namespace + +int main() { + const DraftWeights w = random_model(); + const std::vector prompt = {1, 6, 11, 2}; + const int WANT = 30; + Options o; + o.max_tokens = WANT; + o.tau = 0.0f; + TapBase pb; + const Result plain = onebit::npu::mtp::run_plain(pb, prompt, o); + check(plain.ok && int(plain.tokens.size()) == WANT, "plain decode on the tap base"); + const char* names[] = {"right", "every third wrong", "wrong"}; + for (int mode = 0; mode < 3; mode++) { + for (int k : {2, 7}) { + TapBase b; + int calls = 0; + HeadHooks hooks; + hooks.embed = [](int t, std::vector& r) { embed(t, r); return t >= 0 && t < V; }; + hooks.pick = [&](const float*, int prev) { + const int t = next_of(prev); + const bool wrong = mode == 2 || (mode == 1 && ++calls % 3 == 0); + return wrong ? (t + 1) % V : t; + }; + DFlashHead head(w, b.taps, hooks, /*mask*/ 15, BLOCK); + Logged lg(head); + o.k = k; + const Result r = onebit::npu::mtp::run(b, &lg, prompt, o); + const std::string tag = std::string("k=") + std::to_string(k) + " pick " + names[mode] + ": "; + check(r.ok && r.tokens == plain.tokens && head.error().empty(), tag + "text equals plain decode" + + (head.error().empty() ? "" : " (" + head.error() + ")")); + // [2] the committed stream, rebuilt from scratch + std::vector stream = prompt; + stream.insert(stream.end(), r.tokens.begin(), r.tokens.end()); + std::vector want(stream.size() * size_t(T * H)); + for (size_t p = 0; p < stream.size(); p++) tap_row(stream[p], int(p), want.data() + p * size_t(T * H)); + check(b.taps == want, tag + "the sink holds exactly the " + std::to_string(stream.size()) + + " committed rows (has " + std::to_string(b.taps.size() / size_t(T * H)) + ")"); + // [3] every draft against draft_forward over the committed rows it should have seen + int same = 0; + for (size_t i = 0; i < lg.out.size(); i++) { + std::vector noise(size_t(BLOCK) * H), row, mask; + embed(lg.last[i], row); + embed(15, mask); + std::copy(row.begin(), row.end(), noise.begin()); + for (int j = 1; j < BLOCK; j++) std::copy(mask.begin(), mask.end(), noise.begin() + j * H); + const std::vector ref = + onebit::npu::dflash::draft_forward(w, want.data(), lg.ctx[i], noise.data(), BLOCK, 0); + same += ref == lg.out[i] ? 1 : 0; + } + check(!lg.out.empty() && same == int(lg.out.size()), + tag + std::to_string(same) + " of " + std::to_string(lg.out.size()) + + " drafts equal draft_forward over the committed rows, bit for bit"); + // [4] + if (mode == 0) check(r.stats.drafted > 0 && r.stats.accepted == r.stats.drafted, tag + "every draft accepted"); + if (mode == 1) check(r.stats.accepted > 0 && r.stats.accepted < r.stats.drafted, tag + "a prefix accepted"); + if (mode == 2) check(r.stats.drafted > 0 && r.stats.accepted == 0, tag + "no draft accepted"); + std::printf("INFO %saccepted %d of %d drafts in %d rounds\n", tag.c_str(), r.stats.accepted, + r.stats.drafted, r.stats.iterations); + } + } + { + std::vector taps(size_t(3 * T * H), 0.25f); + HeadHooks hooks; + hooks.embed = [](int t, std::vector& r) { embed(t, r); return true; }; + hooks.pick = [](const float*, int prev) { return next_of(prev); }; + DFlashHead head(w, taps, hooks, 15, BLOCK); + std::vector out; + const int n = head.draft({}, 1, 6, 4, out); + check(n == 0 && out.empty() && head.error().find("position 6 over 3") != std::string::npos, + "a draft at a position the sink does not end at is refused: " + head.error()); + check(head.draft({}, 1, 4, 4, out) == 4 && head.error().empty() && head.last_context_rows() == 3, + "the same head at the position after its 3 rows drafts 4 tokens over 3 context rows"); + } + std::printf(fails ? "\n%d CHECK(S) FAILED\n" : "\nall checks passed\n", fails); + return fails ? 1 : 0; +}