npu: the DFlash draft's forward as a host reference, with hand-computed cases (refs #1780)

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)

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 <dir>`
@@ -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=<dir>` 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=<dir>` 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
 

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 <algorithm>
+#include <cmath>
+#include <cstddef>
+#include <stdexcept>
+#include <string>
+
+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<float> linear(const std::vector<float>& W, int n_out, int n_in, const float* x, int rows) {
+    std::vector<float> 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<float>& 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<float> 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<float> 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<float> x(noise, noise + size_t(block) * size_t(H));
+    std::vector<float> kv_in(size_t(keys) * size_t(H));
+    std::copy(c.begin(), c.end(), kv_in.begin());
+    std::vector<double> score(static_cast<size_t>(keys));
+    std::vector<float> attn(size_t(block) * size_t(Q));
+
+    for (const DraftLayerWeights& L : w.layers) {
+        // attention
+        std::vector<float> 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<float> q = linear(L.q, Q, H, h.data(), block);
+        std::vector<float> k = linear(L.k, KV, H, kv_in.data(), keys);
+        std::vector<float> 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<float> 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<float> p = x;
+        rms_norm(p.data(), size_t(block), H, L.post_norm, w.rms_eps);
+        std::vector<float> g = linear(L.gate, I, H, p.data(), block);
+        std::vector<float> 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<float> 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

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 <vector>
+
+#include "dflash_draft.h"
+
+namespace onebit::npu::dflash {
+
+struct DraftLayerWeights {
+    std::vector<float> q, k, v, o, gate, up, down;                 // row-major [out][in]
+    std::vector<float> 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<float> fc;            // [hidden][taps * hidden]
+    std::vector<float> hidden_norm;   // [hidden]
+    std::vector<float> norm;          // [hidden]
+    std::vector<DraftLayerWeights> 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<float> 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