diff --git a/CMakeLists.txt b/CMakeLists.txt index 3784c0d56..12feff585 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -1360,6 +1360,13 @@ target_link_libraries(dflash_draft_test PRIVATE onebit_npu) add_dependencies(onebit dflash_draft_test) set(ONEBIT_DFLASH_DRAFT "" CACHE PATH "A DFlash draft package directory for dflash_draft_test's real-package arm") add_test(NAME dflash_draft COMMAND dflash_draft_test ${ONEBIT_DFLASH_DRAFT}) +# The draft's forward as a host reference (engine#1780): hand-computed cases and structure properties; the real +# package arm runs when ONEBIT_DFLASH_DRAFT names one. +add_executable(dflash_forward_test tests/dflash_forward_test.cpp npu/dflash_forward.cpp npu/dflash_draft.cpp npu/forward/src/q4nx_dequant.c) +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}) 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/amd-drafters.md b/docs/amd-drafters.md index a7a7235d1..216c5a1f5 100644 --- a/docs/amd-drafters.md +++ b/docs/amd-drafters.md @@ -88,7 +88,7 @@ So the Gumiho row is *"fails as published, because our verifier is the wrong sha quality (its published numbers are the best in its class: 2.97× mean over Eagle-2's 2.74× on vicuna-7B). Running it would mean implementing a tree verifier first, which is its own project. -## DFlash draft packages on the NPU route: **the container is in forms we already read; the forward is what is missing** (engine#1780) +## DFlash draft packages on the NPU route: **the container is in forms we already read; the forward exists as a host reference, not yet on the device** (engine#1780) FastFlowLM 1.0.12 ships speculative decode as two-package tags: a small DFlash draft beside the main model's package. A draft package is `config.json` (`model_type: dflash`) plus one Q4NX container. `tools/dflash_draft_info.py ` @@ -136,7 +136,35 @@ Two facts about the container that the loader depends on, both measured on the B `dflash_draft_test` builds a synthetic package (every projection encoded tile by tile from a formula and compared at every element, the fusion matrix checked for row-major order, sixteen refusals each matched on its message) and needs no model; given a package directory (`-DONEBIT_DFLASH_DRAFT=` or `ONEBIT_TEST_DFLASH_DRAFT`) it also loads and -decodes the real one. The loader is not wired into `forward-serve` yet: there is no forward for it to feed. +decodes the real one. The loader is not wired into `forward-serve` yet. + +### The forward, as a host reference (`npu/dflash_forward.h`) + +`decode_draft(load_draft(dir))` turns a package into f32 weights (about 1.3 GB for the Bolt draft) and +`draft_forward(weights, taps, ctx, noise, block)` runs one draft pass on the host: + +| step | what is computed | +|---|---| +| context | one row per context position: the target's hidden states at `target_layer_ids`, concatenated, through `fc` and then `hidden_norm`. The same rows are used by every layer. | +| block | the target's embeddings of `[last token, mask, mask, ...]`, `block_size` rows. Only these rows change from layer to layer. | +| each layer | `input_norm` on the block rows; queries from the block rows; keys and values from the context rows followed by the block rows; `q_norm` / `k_norm` per head; rotary positions (context first, then the block, consecutive); every block row sees every key, so the block is not causal; then the gated MLP. | +| output | the final `norm` of the block rows. The target's own head turns each row into scores. | + +The draft has no embedding table and no head of its own: both are the target's. + +This is a reference to compare a device route against, not a fast path (plain loops, one thread: a pass of 2 context +and 8 block rows over the Bolt draft takes seconds). **It has not been run against a target model's hidden states +yet**, so nothing here says the draft's proposals are accepted; that needs the target pack's layer taps, which is the +next step of engine#1780. What `dflash_forward_test` establishes: + +- two cases worked out by hand in the test (one block row alone; one context row plus one block row, with a different + gain on every norm that touches the attention), matched to 1e-5; +- the structure, on a small random model: the output does not depend on where the positions start, a later block row + changes an earlier one, the context's scale does not matter, a window of one position equals each row on its own, + and a query head reads the key/value head of its group; +- seven refusals, each matched on its message; +- with a package directory (`-DONEBIT_DFLASH_DRAFT=` or `ONEBIT_TEST_DFLASH_DRAFT`), the real draft decodes and + one pass over a made-up context is finite and position-independent. ## PARD-2's repos are weights-only — borrow the target's tokenizer diff --git a/npu/dflash_forward.cpp b/npu/dflash_forward.cpp new file mode 100644 index 000000000..a9c56a9e6 --- /dev/null +++ b/npu/dflash_forward.cpp @@ -0,0 +1,229 @@ +// 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_forward.cpp -- see dflash_forward.h for what is computed and why it is written this plainly. + +#include "dflash_forward.h" + +#include +#include +#include +#include +#include + +namespace onebit::npu::dflash { + +namespace { + +void need(bool ok, const std::string& what) { + if (!ok) throw std::runtime_error("dflash forward: " + what); +} + +// y[rows][n_out] = x[rows][n_in] * W^T, W row-major [n_out][n_in]. +std::vector linear(const std::vector& W, int n_out, int n_in, const float* x, int rows) { + std::vector y(size_t(rows) * size_t(n_out)); + for (int r = 0; r < rows; r++) { + const float* xr = x + size_t(r) * size_t(n_in); + for (int n = 0; n < n_out; n++) { + const float* wr = W.data() + size_t(n) * size_t(n_in); + double acc = 0.0; + for (int k = 0; k < n_in; k++) acc += double(wr[k]) * double(xr[k]); + y[size_t(r) * size_t(n_out) + size_t(n)] = float(acc); + } + } + return y; +} + +// RMS norm over groups of `width` values: x * gain / sqrt(mean(x^2) + eps). `gain` has `width` entries. +void rms_norm(float* x, size_t groups, int width, const std::vector& gain, float eps) { + for (size_t g = 0; g < groups; g++) { + float* v = x + g * size_t(width); + double ss = 0.0; + for (int i = 0; i < width; i++) ss += double(v[i]) * double(v[i]); + const double inv = 1.0 / std::sqrt(ss / double(width) + double(eps)); + for (int i = 0; i < width; i++) v[i] = float(double(v[i]) * inv * double(gain[size_t(i)])); + } +} + +// Rotate one head vector of width d at position `pos`: pairs (i, i + d/2) turn by pos * theta^(-2i/d). +void rope(float* v, int d, long pos, float theta) { + const int half = d / 2; + for (int i = 0; i < half; i++) { + const double ang = double(pos) * std::pow(double(theta), -2.0 * double(i) / double(d)); + const double c = std::cos(ang), s = std::sin(ang); + const double a = v[i], b = v[i + half]; + v[i] = float(a * c - b * s); + v[i + half] = float(b * c + a * s); + } +} + +} // namespace + +void check_weights(const DraftWeights& w) { + need(w.hidden > 0 && w.intermediate > 0 && w.heads > 0 && w.kv_heads > 0 && w.head_dim > 0 && w.taps > 0, + "a geometry field is zero or negative"); + need(w.head_dim % 2 == 0, "head_dim " + std::to_string(w.head_dim) + " is odd, so it cannot be rotated in pairs"); + need(w.heads % w.kv_heads == 0, std::to_string(w.heads) + " heads are not a multiple of " + + std::to_string(w.kv_heads) + " key/value heads"); + need(w.rope_theta > 0.0f, "rope_theta is not positive"); + need(!w.layers.empty(), "no layers"); + const size_t H = size_t(w.hidden), I = size_t(w.intermediate); + const size_t Q = size_t(w.heads) * size_t(w.head_dim), KV = size_t(w.kv_heads) * size_t(w.head_dim); + need(w.fc.size() == H * H * size_t(w.taps), "fc is not hidden x (taps * hidden)"); + need(w.hidden_norm.size() == H && w.norm.size() == H, "hidden_norm or norm is not hidden wide"); + for (size_t l = 0; l < w.layers.size(); l++) { + const DraftLayerWeights& L = w.layers[l]; + const std::string at = "layer " + std::to_string(l) + ": "; + need(L.q.size() == Q * H, at + "q_proj is not (heads * head_dim) x hidden"); + need(L.k.size() == KV * H && L.v.size() == KV * H, at + "k_proj or v_proj is not (kv_heads * head_dim) x hidden"); + need(L.o.size() == H * Q, at + "o_proj is not hidden x (heads * head_dim)"); + need(L.gate.size() == I * H && L.up.size() == I * H, at + "gate or up is not intermediate x hidden"); + need(L.down.size() == H * I, at + "down is not hidden x intermediate"); + need(L.input_norm.size() == H && L.post_norm.size() == H, at + "a layer norm is not hidden wide"); + need(L.q_norm.size() == size_t(w.head_dim) && L.k_norm.size() == size_t(w.head_dim), + at + "q_norm or k_norm is not head_dim wide"); + } +} + +DraftWeights decode_draft(const DraftModel& m) { + DraftWeights w; + w.hidden = m.cfg.hidden; + w.intermediate = m.cfg.intermediate; + w.heads = m.cfg.heads; + w.kv_heads = m.cfg.kv_heads; + w.head_dim = m.cfg.head_dim; + w.taps = m.cfg.taps(); + w.rms_eps = m.cfg.rms_eps; + w.rope_theta = m.cfg.rope_theta; + w.fc = decode_fusion(m.fc); + w.hidden_norm = decode_gain(m.hidden_norm); + w.norm = decode_gain(m.norm); + w.layers.reserve(m.layers.size()); + for (const DraftLayer& L : m.layers) { + DraftLayerWeights d; + d.q = decode_projection(L.q); + d.k = decode_projection(L.k); + d.v = decode_projection(L.v); + d.o = decode_projection(L.o); + d.gate = decode_projection(L.gate); + d.up = decode_projection(L.up); + d.down = decode_projection(L.down); + d.input_norm = decode_gain(L.input_norm); + d.post_norm = decode_gain(L.post_norm); + d.q_norm = decode_gain(L.q_norm); + d.k_norm = decode_gain(L.k_norm); + w.layers.push_back(std::move(d)); + } + check_weights(w); + return w; +} + +std::vector draft_forward(const DraftWeights& w, const float* taps, int ctx, const float* noise, int block, + long pos0, int window) { + check_weights(w); + need(block > 0 && noise != nullptr, "an empty block"); + need(ctx >= 0 && (ctx == 0 || taps != nullptr), "a context length without context rows"); + need(window >= 0, "a negative window"); + const int H = w.hidden, I = w.intermediate, D = w.head_dim; + const int Q = w.heads * D, KV = w.kv_heads * D; + const int keys = ctx + block; + const int group = w.heads / w.kv_heads; // query heads per key/value head, consecutive + + // The context: fused, normalised once, and the same rows for every layer. + std::vector c; + if (ctx > 0) { + c = linear(w.fc, H, H * w.taps, taps, ctx); + rms_norm(c.data(), size_t(ctx), H, w.hidden_norm, w.rms_eps); + } + + std::vector x(noise, noise + size_t(block) * size_t(H)); + std::vector kv_in(size_t(keys) * size_t(H)); + std::copy(c.begin(), c.end(), kv_in.begin()); + std::vector score(static_cast(keys)); + std::vector attn(size_t(block) * size_t(Q)); + + for (const DraftLayerWeights& L : w.layers) { + // attention + std::vector h = x; + rms_norm(h.data(), size_t(block), H, L.input_norm, w.rms_eps); + std::copy(h.begin(), h.end(), kv_in.begin() + ptrdiff_t(size_t(ctx) * size_t(H))); + + std::vector q = linear(L.q, Q, H, h.data(), block); + std::vector k = linear(L.k, KV, H, kv_in.data(), keys); + std::vector v = linear(L.v, KV, H, kv_in.data(), keys); + rms_norm(q.data(), size_t(block) * size_t(w.heads), D, L.q_norm, w.rms_eps); + rms_norm(k.data(), size_t(keys) * size_t(w.kv_heads), D, L.k_norm, w.rms_eps); + for (int i = 0; i < block; i++) + for (int hd = 0; hd < w.heads; hd++) + rope(q.data() + size_t(i) * size_t(Q) + size_t(hd) * size_t(D), D, pos0 + ctx + i, w.rope_theta); + for (int t = 0; t < keys; t++) + for (int hd = 0; hd < w.kv_heads; hd++) + rope(k.data() + size_t(t) * size_t(KV) + size_t(hd) * size_t(D), D, pos0 + t, w.rope_theta); + + const double scale = 1.0 / std::sqrt(double(D)); + for (int i = 0; i < block; i++) { + for (int hd = 0; hd < w.heads; hd++) { + const int kh = hd / group; + const float* qv = q.data() + size_t(i) * size_t(Q) + size_t(hd) * size_t(D); + double top = -INFINITY; + for (int t = 0; t < keys; t++) { + const long dist = long(ctx + i) - long(t); + if (window > 0 && (dist >= window || -dist >= window)) { + score[size_t(t)] = -INFINITY; + continue; + } + const float* kvp = k.data() + size_t(t) * size_t(KV) + size_t(kh) * size_t(D); + double s = 0.0; + for (int d = 0; d < D; d++) s += double(qv[d]) * double(kvp[d]); + score[size_t(t)] = s * scale; + if (score[size_t(t)] > top) top = score[size_t(t)]; + } + double sum = 0.0; + for (int t = 0; t < keys; t++) { + score[size_t(t)] = std::isinf(score[size_t(t)]) ? 0.0 : std::exp(score[size_t(t)] - top); + sum += score[size_t(t)]; + } + need(sum > 0.0 && std::isfinite(sum), "an attention row has no visible key or a non-finite score"); + float* out = attn.data() + size_t(i) * size_t(Q) + size_t(hd) * size_t(D); + for (int d = 0; d < D; d++) { + double acc = 0.0; + for (int t = 0; t < keys; t++) + acc += score[size_t(t)] * double(v[size_t(t) * size_t(KV) + size_t(kh) * size_t(D) + size_t(d)]); + out[d] = float(acc / sum); + } + } + } + std::vector o = linear(L.o, H, Q, attn.data(), block); + for (size_t j = 0; j < x.size(); j++) x[j] += o[j]; + + // MLP + std::vector p = x; + rms_norm(p.data(), size_t(block), H, L.post_norm, w.rms_eps); + std::vector g = linear(L.gate, I, H, p.data(), block); + std::vector u = linear(L.up, I, H, p.data(), block); + for (size_t j = 0; j < g.size(); j++) { + const double a = g[j]; + g[j] = float(a / (1.0 + std::exp(-a)) * double(u[j])); // silu(gate) * up + } + std::vector dn = linear(L.down, H, I, g.data(), block); + for (size_t j = 0; j < x.size(); j++) x[j] += dn[j]; + } + + rms_norm(x.data(), size_t(block), H, w.norm, w.rms_eps); + for (float f : x) need(std::isfinite(f), "a non-finite output"); + return x; +} + +} // namespace onebit::npu::dflash diff --git a/npu/dflash_forward.h b/npu/dflash_forward.h new file mode 100644 index 000000000..d739c85c9 --- /dev/null +++ b/npu/dflash_forward.h @@ -0,0 +1,78 @@ +// 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_forward.h -- the DFlash draft's forward as a host reference in f32 (engine#1780, item 2 of "supported"). +// +// What a draft computes, from the interface its package states (dflash_draft.h loads it): +// +// context c[t] = hidden_norm(fc(taps[t])) one row per context position; taps[t] is the target +// model's hidden states at target_layer_ids, concatenated +// block x[i] = noise[i] the target's embeddings of [last token, mask, mask, ...] +// per layer h = input_norm(x) +// q = q_norm(q_proj(h)) the block rows only +// k, v = k_norm(k_proj([c ; h])), v_proj([c ; h]) context rows first, then the block rows +// x += o_proj(attention(rope(q), rope(k), v)) every block row sees every key: not causal +// x += down(silu(gate(p)) * up(p)), p = post_norm(x) +// output norm(x) [block][hidden]; the target's own head scores it +// +// The context rows are NOT passed through a layer's input_norm and are the same rows in every layer; only the block +// rows change from layer to layer. Positions are consecutive: context position t is pos0 + t and block row i is +// pos0 + ctx + i. The rotation is relative, so the output does not depend on pos0 (the test holds it to that). +// +// This is the reference the device route will be compared with, not a fast path: plain loops, one thread, f32 sums +// accumulated in double. It has not yet been run against a target model's own hidden states (that needs the target +// pack's layer taps, the next step), so what the test establishes is the structure above and a hand-computed case, +// not agreement with a published draft's outputs. + +#pragma once + +#include + +#include "dflash_draft.h" + +namespace onebit::npu::dflash { + +struct DraftLayerWeights { + std::vector q, k, v, o, gate, up, down; // row-major [out][in] + std::vector input_norm, post_norm, q_norm, k_norm; // gains +}; + +// A draft decoded to f32. For the Bolt draft this is about 1.3 GB; the caller decides how long to keep it. +struct DraftWeights { + int hidden = 0, intermediate = 0, heads = 0, kv_heads = 0, head_dim = 0, taps = 0; + float rms_eps = 0.0f, rope_theta = 0.0f; + std::vector fc; // [hidden][taps * hidden] + std::vector hidden_norm; // [hidden] + std::vector norm; // [hidden] + std::vector layers; +}; + +// Decode every tensor of a loaded package. Throws std::runtime_error if a decoded tensor has the wrong size. +DraftWeights decode_draft(const DraftModel& m); + +// Throws std::runtime_error naming the first thing wrong with a DraftWeights' geometry or sizes. +void check_weights(const DraftWeights& w); + +// One draft pass. +// taps [ctx][taps * hidden], row t = the target's hidden states at position pos0 + t, in target_layer_ids order +// noise [block][hidden] +// window 0 = every key is visible; otherwise a block row sees a key only when their positions differ by less +// than `window` (both directions, as the draft is not causal) +// Returns [block][hidden]. ctx may be 0 (the block then attends only to itself). Throws std::runtime_error on a bad +// argument or a non-finite result. +std::vector draft_forward(const DraftWeights& w, const float* taps, int ctx, const float* noise, int block, + long pos0 = 0, int window = 0); + +} // namespace onebit::npu::dflash diff --git a/tests/dflash_forward_test.cpp b/tests/dflash_forward_test.cpp new file mode 100644 index 000000000..94f0b8128 --- /dev/null +++ b/tests/dflash_forward_test.cpp @@ -0,0 +1,340 @@ +// 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_forward_test.cpp -- the DFlash draft's host forward (npu/dflash_forward.h, engine#1780). +// +// Three kinds of arm, none needing a device: +// 1. two cases small enough to work out by hand, written out here as plain arithmetic: one block row with no +// context (norm, value path, residuals, MLP, final norm), and one block row with one context row (the fused and +// normalised context, the query/key norms, the rotation and the two-key softmax); +// 2. properties of the structure on a small random model: the output does not depend on where the positions start, +// a later block row changes an earlier one (not causal), the context is read and its scale does not matter, a +// window of one position is the same as each row alone, and a query head reads the key/value head of its group; +// 3. the real Bolt draft, when a path is given (argv[1] or ONEBIT_TEST_DFLASH_DRAFT): it decodes and one pass over +// a made-up context is finite and position-independent. That arm does not show the draft drafts well -- it has +// no target model -- only that the real weights go through this code. + +#include +#include +#include +#include +#include +#include +#include + +#include "dflash_forward.h" + +namespace { + +using namespace onebit::npu::dflash; + +int g_failures = 0; +int g_checks = 0; + +void check(bool cond, const std::string& what) { + ++g_checks; + if (!cond) { + std::printf("FAIL %s\n", what.c_str()); + ++g_failures; + } +} + +double max_diff(const std::vector& a, const std::vector& b) { + if (a.size() != b.size()) return 1e30; + double m = 0.0; + for (size_t i = 0; i < a.size(); i++) m = std::max(m, std::fabs(double(a[i]) - double(b[i]))); + return m; +} + +bool throws(const std::function& f, const char* needle) { + try { + f(); + } catch (const std::runtime_error& e) { + return std::string(e.what()).find(needle) != std::string::npos; + } + return false; +} + +struct Rng { + uint64_t s; + float next() { // uniform in [-1, 1) + s = s * 6364136223846793005ULL + 1442695040888963407ULL; + return float(double((s >> 40) & 0xFFFFFF) / double(0x800000) - 1.0); + } + std::vector vec(size_t n, float scale) { + std::vector v(n); + for (float& f : v) f = next() * scale; + return v; + } +}; + +std::vector identity(int n) { + std::vector m(size_t(n) * size_t(n), 0.0f); + for (int i = 0; i < n; i++) m[size_t(i) * size_t(n) + size_t(i)] = 1.0f; + return m; +} + +// hidden 2, one head of width 2, one layer, eps 0, every projection the identity except v, gate and up. The gains that sit on +// the attention's inputs are all different, so a norm applied to the wrong rows (or not at all) moves the answer. +DraftWeights tiny_model() { + DraftWeights w; + w.hidden = 2; w.intermediate = 2; w.heads = 1; w.kv_heads = 1; w.head_dim = 2; w.taps = 1; + w.rms_eps = 0.0f; w.rope_theta = 10000.0f; + w.fc = identity(2); + w.hidden_norm = {0.5f, 3.0f}; + w.norm = {1.0f, 1.0f}; + DraftLayerWeights L; + L.q = L.k = L.o = L.down = identity(2); + L.v = {2.0f, 0.0f, 0.0f, 3.0f}; + // gate and up differ, and gate drives its first output negative where up is positive, so silu(gate) * up and + // silu(up) * gate are different numbers: which of the two is activated is part of the answer. + L.gate = {-1.0f, 0.0f, 0.0f, 2.0f}; + L.up = {3.0f, 0.0f, 0.0f, 0.5f}; + L.input_norm = {2.0f, 0.5f}; + L.q_norm = {1.5f, 0.5f}; + L.k_norm = {0.25f, 2.0f}; + L.post_norm = {1.0f, 1.0f}; + w.layers.push_back(L); + return w; +} + +void rms2(double& a, double& b) { + const double r = std::sqrt((a * a + b * b) / 2.0); + a /= r; b /= r; +} +void rmsg(double& a, double& b, double g0, double g1) { // RMS norm, then the two gains + rms2(a, b); + a *= g0; b *= g1; +} +double silu(double a) { return a / (1.0 + std::exp(-a)); } + +// What follows the attention in the tiny model, given the block row and the attention's output. +void tiny_tail(double x0, double x1, double a0, double a1, double& o0, double& o1) { + double y0 = x0 + a0, y1 = x1 + a1; // residual (o_proj is the identity) + double p0 = y0, p1 = y1; + rms2(p0, p1); // post_norm + y0 += silu(-1.0 * p0) * (3.0 * p0); // silu(gate(p)) * up(p); down is the identity + y1 += silu(2.0 * p1) * (0.5 * p1); + rms2(y0, y1); // the final norm + o0 = y0; o1 = y1; +} + +void hand_cases() { + const DraftWeights w = tiny_model(); + { // one block row, no context: the only key is the row itself, so the attention returns v(h) exactly + const float noise[2] = {3.0f, 4.0f}; + double h0 = 3.0, h1 = 4.0; + rmsg(h0, h1, 2.0, 0.5); // input_norm + double e0, e1; + tiny_tail(3.0, 4.0, 2.0 * h0, 3.0 * h1, e0, e1); + const std::vector out = draft_forward(w, nullptr, 0, noise, 1); + check(out.size() == 2 && std::fabs(out[0] - e0) < 1e-5 && std::fabs(out[1] - e1) < 1e-5, + "one row, no context: the output is the hand-computed one"); + std::printf("hand case 1: got %.6f %.6f, expected %.6f %.6f\n", out[0], out[1], e0, e1); + } + { // one context row and one block row: two keys, at positions 5 and 6 + const float tap[2] = {1.0f, 0.5f}; + const float noise[2] = {-2.0f, 1.0f}; + double c0 = 1.0, c1 = 0.5; + rmsg(c0, c1, 0.5, 3.0); // fc is the identity, then hidden_norm; no input_norm on this row + double h0 = -2.0, h1 = 1.0; + rmsg(h0, h1, 2.0, 0.5); // input_norm, the block row only + double qn0 = h0, qn1 = h1, kcn0 = c0, kcn1 = c1, khn0 = h0, khn1 = h1; + rmsg(qn0, qn1, 1.5, 0.5); // q_norm + rmsg(kcn0, kcn1, 0.25, 2.0); // k_norm, on both keys + rmsg(khn0, khn1, 0.25, 2.0); + auto rot = [](double a, double b, double ang, double& ra, double& rb) { + ra = a * std::cos(ang) - b * std::sin(ang); + rb = b * std::cos(ang) + a * std::sin(ang); + }; + double q0, q1, kc0, kc1, kh0, kh1; + rot(qn0, qn1, 6.0, q0, q1); // head_dim 2: one pair, turned by its position + rot(kcn0, kcn1, 5.0, kc0, kc1); + rot(khn0, khn1, 6.0, kh0, kh1); + const double sc = (q0 * kc0 + q1 * kc1) / std::sqrt(2.0), sh = (q0 * kh0 + q1 * kh1) / std::sqrt(2.0); + const double wc = std::exp(sc) / (std::exp(sc) + std::exp(sh)), wh = 1.0 - wc; + const double a0 = wc * 2.0 * c0 + wh * 2.0 * h0, a1 = wc * 3.0 * c1 + wh * 3.0 * h1; // v = diag(2, 3) + double e0, e1; + tiny_tail(-2.0, 1.0, a0, a1, e0, e1); + const std::vector out = draft_forward(w, tap, 1, noise, 1, 5); + check(out.size() == 2 && std::fabs(out[0] - e0) < 1e-5 && std::fabs(out[1] - e1) < 1e-5, + "one context row and one block row: the output is the hand-computed one"); + std::printf("hand case 2: got %.6f %.6f, expected %.6f %.6f (context weight %.4f)\n", out[0], out[1], e0, e1, wc); + } +} + +DraftWeights random_model(uint64_t seed) { + Rng r{seed}; + DraftWeights w; + w.hidden = 8; w.intermediate = 12; w.heads = 4; w.kv_heads = 2; w.head_dim = 4; w.taps = 2; + w.rms_eps = 1e-6f; w.rope_theta = 10000.0f; + const size_t H = 8, I = 12, Q = 16, KV = 8; + w.fc = r.vec(H * H * 2, 0.5f); + auto gain = [&](size_t n) { + std::vector g = r.vec(n, 0.3f); + for (float& f : g) f += 1.0f; + return g; + }; + w.hidden_norm = gain(H); + w.norm = gain(H); + for (int l = 0; l < 2; l++) { + DraftLayerWeights L; + L.q = r.vec(Q * H, 0.6f); L.k = r.vec(KV * H, 0.6f); L.v = r.vec(KV * H, 0.6f); L.o = r.vec(H * Q, 0.4f); + L.gate = r.vec(I * H, 0.5f); L.up = r.vec(I * H, 0.5f); L.down = r.vec(H * I, 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 properties() { + const DraftWeights w = random_model(7); + Rng r{99}; + const int ctx = 3, block = 4, H = 8; + const std::vector taps = r.vec(size_t(ctx) * 16, 1.0f); + const std::vector noise = r.vec(size_t(block) * size_t(H), 1.0f); + const std::vector base = draft_forward(w, taps.data(), ctx, noise.data(), block); + check(base.size() == size_t(block) * size_t(H), "the output is block x hidden"); + + check(max_diff(base, draft_forward(w, taps.data(), ctx, noise.data(), block, 1000)) < 2e-4, + "the output does not depend on where the positions start"); + + { + std::vector n2 = noise; + n2[size_t(block - 1) * size_t(H) + 2] += 0.5f; // the last row + const std::vector o = draft_forward(w, taps.data(), ctx, n2.data(), block); + double first_row = 0.0; + for (int j = 0; j < H; j++) first_row = std::max(first_row, std::fabs(double(o[size_t(j)]) - double(base[size_t(j)]))); + check(first_row > 1e-4, "a change in the last block row reaches the first one: the draft is not causal"); + } + { + std::vector t2 = taps; + t2[5] += 0.5f; + check(max_diff(base, draft_forward(w, t2.data(), ctx, noise.data(), block)) > 1e-4, "the context is read"); + check(max_diff(base, draft_forward(w, nullptr, 0, noise.data(), block)) > 1e-4, + "a pass without context differs from one with it"); + std::vector t7 = taps; + for (float& f : t7) f *= 7.0f; + check(max_diff(base, draft_forward(w, t7.data(), ctx, noise.data(), block)) < 2e-4, + "the context's scale does not matter: it is normalised after fusion, before every layer"); + } + { + check(max_diff(base, draft_forward(w, taps.data(), ctx, noise.data(), block, 0, ctx + block)) == 0.0, + "a window as wide as the pass hides nothing"); + const std::vector w1 = draft_forward(w, taps.data(), ctx, noise.data(), block, 0, 1); + std::vector alone; + for (int i = 0; i < block; i++) { + const std::vector one = draft_forward(w, nullptr, 0, noise.data() + size_t(i) * size_t(H), 1); + alone.insert(alone.end(), one.begin(), one.end()); + } + check(max_diff(w1, alone) < 2e-4, "a window of one position is each block row on its own"); + check(max_diff(w1, base) > 1e-4, "and that differs from the unwindowed pass"); + } + { // query heads 0,1 read key/value head 0 and heads 2,3 read head 1. Keep only head 2's columns of o_proj + // (in a one-layer model), then the output must ignore key/value head 0 and follow head 1. + DraftWeights g = random_model(11); + g.layers.resize(1); + const int D = 4, Q = 16; + for (int n = 0; n < H; n++) + for (int col = 0; col < Q; col++) + if (col / D != 2) g.layers[0].o[size_t(n) * size_t(Q) + size_t(col)] = 0.0f; + const std::vector ref = draft_forward(g, taps.data(), ctx, noise.data(), block); + DraftWeights g0 = g, g1 = g; + for (int row = 0; row < D; row++) + for (int col = 0; col < H; col++) { + g0.layers[0].v[size_t(row) * size_t(H) + size_t(col)] += 0.25f; // value head 0 + g1.layers[0].v[size_t(row + D) * size_t(H) + size_t(col)] += 0.25f; // value head 1 + } + check(max_diff(ref, draft_forward(g0, taps.data(), ctx, noise.data(), block)) == 0.0, + "query head 2 does not read key/value head 0"); + check(max_diff(ref, draft_forward(g1, taps.data(), ctx, noise.data(), block)) > 1e-4, + "query head 2 reads key/value head 1"); + } +} + +void refusals() { + const DraftWeights w = random_model(3); + const std::vector noise(8, 0.5f); + check(throws([&] { draft_forward(w, nullptr, 0, noise.data(), 0); }, "empty block"), "an empty block is refused"); + check(throws([&] { draft_forward(w, nullptr, 2, noise.data(), 1); }, "without context rows"), + "a context length with no rows is refused"); + check(throws([&] { draft_forward(w, nullptr, 0, noise.data(), 1, 0, -1); }, "negative window"), + "a negative window is refused"); + { + DraftWeights b = w; + b.layers[1].k.pop_back(); + check(throws([&] { draft_forward(b, nullptr, 0, noise.data(), 1); }, "layer 1: k_proj"), + "a projection of the wrong size is refused with its layer"); + } + { + DraftWeights b = w; + b.kv_heads = 3; + check(throws([&] { check_weights(b); }, "not a multiple"), "heads that do not divide into key/value heads are refused"); + } + { + DraftWeights b = w; + b.head_dim = 3; + check(throws([&] { check_weights(b); }, "odd"), "an odd head width is refused"); + } + { + std::vector bad = noise; + bad[3] = NAN; + check(throws([&] { draft_forward(w, nullptr, 0, bad.data(), 1); }, "non-finite"), "a non-finite result is refused"); + } +} + +void real_package(const std::string& dir) { + DraftWeights w; + try { + w = decode_draft(load_draft(dir)); + } catch (const std::exception& e) { + check(false, std::string("the real package decodes: ") + e.what()); + return; + } + const int ctx = 2, block = 8; + Rng r{2026}; + const std::vector taps = r.vec(size_t(ctx) * size_t(w.taps) * size_t(w.hidden), 1.0f); + const std::vector noise = r.vec(size_t(block) * size_t(w.hidden), 1.0f); + try { + const std::vector a = draft_forward(w, taps.data(), ctx, noise.data(), block); + const std::vector b = draft_forward(w, taps.data(), ctx, noise.data(), block, 4096); + double peak = 0.0; + for (float f : a) peak = std::max(peak, std::fabs(double(f))); + check(a.size() == size_t(block) * size_t(w.hidden), "the real package: one pass returns block x hidden"); + check(max_diff(a, b) < 1e-3 * std::max(1.0, peak), "the real package: the pass does not depend on where the positions start"); + std::printf("real package: %zu layers, hidden %d; one pass of %d context + %d block rows is finite, peak |x| %.3f, " + "moved by at most %.2e when the positions start at 4096\n", + w.layers.size(), w.hidden, ctx, block, peak, max_diff(a, b)); + } catch (const std::exception& e) { + check(false, std::string("the real package runs: ") + e.what()); + } +} + +} // namespace + +int main(int argc, char** argv) { + hand_cases(); + properties(); + refusals(); + const char* real = (argc > 1 && argv[1][0] != '\0') ? argv[1] : std::getenv("ONEBIT_TEST_DFLASH_DRAFT"); + if (real && real[0] != '\0') real_package(real); + else std::printf("SKIP the real package: pass its directory, or set ONEBIT_TEST_DFLASH_DRAFT\n"); + if (g_failures) { + std::printf("dflash_forward: %d of %d checks FAILED\n", g_failures, g_checks); + return 1; + } + std::printf("dflash_forward: ok (%d checks)\n", g_checks); + return 0; +}