From 824f941df92ec20e0a973e11cc540e4f1759f239 Mon Sep 17 00:00:00 2001 From: pi-demo Date: Mon, 9 Nov 2026 22:12:20 -0400 Subject: [PATCH 01/20] npu-pack: read a component's per-layer embedding table, beside its token table (WIP, engine#1725) The gemma-4 1.9 components carry, beside `embed_tokens.weight`, a PER-LAYER table and a projection norm: embedding.embed_tokens_per_layer.weight (262144, 8960) f32, external (= 35 layers x 256 head_dim) embedding.per_layer_projection_norm.weight (256,) f32 The decoder consumes them as `per_layer_inputs`, so a container with only the token table gives the decoder less than it asks for -- which is why the previous commit refused such a component by name. This adds the reader, so the next slice can write the container section the forward path will load. EmbedPerLayer { table [V][layers*head_dim] row-major, norm [width], V, width } + `ok()`; read_embedding_per_layer(pack_dir, why) resolves the component exactly as the token half does -- genai_config.json's model.embedding.filename (relative only, no ".."), external data beside the GRAPH rather than at the pack root -- reads f32, refuses any other dtype by name rather than mis-decoding it, and says so when the projection norm is missing (the table is still usable, but a consumer expecting a projected row would silently take raw ones). WIP: no container section and no forward consumer yet; that is the next slice, with its test. Builds clean (`npu_pack_import_test` rc 0). No vendor code or layout: only the pack's own ONNX names and dims are read. --- npu/pack_import.cpp | 79 +++++++++++++++++++++++++++++++++++++++++++++ npu/pack_import.h | 13 ++++++++ 2 files changed, 92 insertions(+) diff --git a/npu/pack_import.cpp b/npu/pack_import.cpp index 19c171787..859a6f723 100644 --- a/npu/pack_import.cpp +++ b/npu/pack_import.cpp @@ -912,6 +912,85 @@ DeclaredMap read_declared(const std::string& pack_dir) { // Returns [V, H] f32 (row-major), empty when the pack has no such component. Everything it needs already existed: // walk_graph() reads the graph with its external-data declarations, decode_matmulnbits() is block-parameterised, and // the scale-per-32-column-group layout is the container's own. +// engine#1725: the per-layer half of a multi-component embedding. Same component path resolution as the token table +// (genai_config.json's model.embedding.filename), same external-data rule (resolved beside the graph), and the same +// dtype rule slice 1 established: f32 is read, anything else unquantised is refused by dtype rather than mis-read. +// Kept separate from read_embedding_component() because the shapes differ (V x width and a width-wide norm, not V x H) +// and because the consumer differs: the token table is the head's ORACLE, these are a decoder INPUT. +EmbedPerLayer read_embedding_per_layer(const std::string& pack_dir, std::string* why) { + EmbedPerLayer out; + std::string gpath = pack_dir + "/embedding.onnx"; + { + std::ifstream gf(pack_dir + "/genai_config.json"); + if (gf) { + json g; + try { g = json::parse(gf); } catch (...) { g = json::object(); } + if (g.is_object() && g.contains("model") && g["model"].is_object() && g["model"].contains("embedding") && + g["model"]["embedding"].is_object()) { + const json& emb = g["model"]["embedding"]; + if (emb.contains("filename") && emb["filename"].is_string()) { + const std::string rel = emb["filename"].get(); + if (!rel.empty() && rel.front() != '/' && rel.find("..") == std::string::npos) { + std::error_code nec; + if (fs::is_regular_file(pack_dir + "/" + rel, nec)) gpath = pack_dir + "/" + rel; + } + } + } + } + } + std::error_code ec; + if (!fs::is_regular_file(gpath, ec)) { if (why) *why = "no embedding component graph at " + gpath; return out; } + OnnxModel em; + try { em = parse_onnx(gpath); } catch (const std::exception& e) { if (why) *why = e.what(); return out; } + const OnnxTensor* table = nullptr; + const OnnxTensor* norm = nullptr; + for (const auto& t : em.tensors) { + if (t.name.find("embed_tokens_per_layer") != std::string::npos) table = &t; + else if (norm == nullptr && t.name.find("per_layer_projection_norm") != std::string::npos) norm = &t; + } + if (table == nullptr) { if (why) *why = "the component carries no embed_tokens_per_layer table"; return out; } + if (table->dims.size() != 2 || table->dims.at(0) <= 0 || table->dims.at(1) <= 0) { + if (why) *why = "embed_tokens_per_layer is not 2-D with positive dims"; return out; + } + if (table->type != kF32) { + if (why) *why = "embed_tokens_per_layer has dtype " + std::to_string(table->type) + ", and only f32 is read here"; + return out; + } + const int64_t V = table->dims.at(0), width = table->dims.at(1); + const auto read_f32 = [&](const OnnxTensor& t, size_t count, const char* what, std::vector* dst) { + const std::string base = fs::path(gpath).parent_path().string() + "/" + t.location; + size_t sz = 0; + std::string w2; + const uint8_t* m = map_file(base, sz, &w2); + if (!m) { if (why) *why = std::string(what) + ": " + w2; return false; } + const uint64_t need = uint64_t(count) * 4ull; + if (t.offset + need > sz) { if (why) *why = std::string(what) + " declares more bytes than its data carries"; return false; } + dst->resize(count); + std::memcpy(dst->data(), m + t.offset, size_t(need)); + return true; + }; + if (!table->location.empty()) { + if (!read_f32(*table, size_t(V) * size_t(width), "embed_tokens_per_layer", &out.table)) return out; + } else if (table->raw != nullptr && table->raw_len >= size_t(V) * size_t(width) * 4) { + out.table.resize(size_t(V) * size_t(width)); + std::memcpy(out.table.data(), table->raw, out.table.size() * sizeof(float)); + } else { + if (why) *why = "embed_tokens_per_layer has no readable data"; return out; + } + if (norm != nullptr && norm->type == kF32 && norm->dims.size() == 1 && norm->dims.at(0) == width && + !norm->location.empty()) { + if (!read_f32(*norm, size_t(width), "per_layer_projection_norm", &out.norm)) return out; + } + if (out.norm.empty()) { + // No projection norm: the table is still usable, but say so — a consumer that expects the projection to be + // normalised would otherwise silently use raw rows. + if (why) *why = "no usable per_layer_projection_norm beside the table (rows are NOT projected)"; + } + out.V = V; + out.width = width; + return out; +} + std::vector read_embedding_component(const std::string& pack_dir, std::string* why) { // A multi-component pack (`type: compositemodel`) names its embedding component in genai_config.json // (`model.embedding.filename = "embedding/model.onnx"`), exactly as it names its decoder; a single-graph pack keeps diff --git a/npu/pack_import.h b/npu/pack_import.h index 1a4f96da4..5c819e357 100644 --- a/npu/pack_import.h +++ b/npu/pack_import.h @@ -78,6 +78,19 @@ struct Dims { // reason. Exposed so the dequantisation can be tested against known codes and scales instead of only on a 3 GB pack. std::vector read_embedding_component(const std::string& pack_dir, std::string* why); +// The per-layer half of a multi-component pack's embedding (engine#1725). gemma-4 components carry, beside the token +// table, `embed_tokens_per_layer.weight` (262144 x 8960 on e2b = layers x head_dim) and a 256-wide projection norm; +// the decoder consumes them as `per_layer_inputs`, so a container without them gives the decoder less than it asks for. +// Read separately from the token table because they have different shapes and different consumers: the token table +// becomes the head oracle, this becomes an input. Empty (with `why`) when the component has no per-layer table. +struct EmbedPerLayer { + std::vector table; // [V][layers*head_dim], row-major + std::vector norm; // [layers*head_dim] + int64_t V = 0, width = 0; + bool ok() const { return V > 0 && width > 0 && table.size() == size_t(V) * size_t(width) && norm.size() == size_t(width); } +}; +EmbedPerLayer read_embedding_per_layer(const std::string& pack_dir, std::string* why); + // One tensor the pack describes, and where its values would come from. struct Planned { std::string name; // Q4NX tensor name (model.layers.N.<...>.weight) From e0001a26e81fe998fcc8838d9eac7be6dc1e9da2 Mon Sep 17 00:00:00 2001 From: pi-demo Date: Mon, 9 Nov 2026 22:15:33 -0400 Subject: [PATCH 02/20] npu-pack: write a component's per-layer embedding into the container (WIP, engine#1725) Second half of the reader in the previous commit: `import_pack()` now writes the per-layer table and its projection norm beside the token table -- model.embed_tokens_per_layer.weight BF16 {V, layers*head_dim} (from the component's f32) model.per_layer_projection_norm.weight BF16 {layers*head_dim} -- and the slice-1 refusal in `read_embedding_component()` is relaxed to match: a component that carries `embed_tokens_per_layer` is no longer refused by that reader, because the mechanism it was waiting for now exists. Measured on the real pack with this change (`1bit npu-pack info ~/scratch-packs/pi-demo/e2b`, loader built from main): before: the pack has no model.embed_tokens.weight, and the component carries embedding.embed_tokens_per_layer.weight, i.e. a PER-LAYER embedding table ... the per-layer embedding mechanism is not implemented here yet after: rc 0, a full plan -- dims hidden 1536 / 35L / 8 heads / 1 kv / head_dim 256 / IM 6144 / vocab 262144, 423 tensors, and 350 `incomplete` entries that all say the same NEXT thing: "layer blocks not located: the epsilon headers do not repeat with a single period" So the embedding half is no longer what blocks a gemma-4 pack: the next wall is the Token .fconst block layout, whose rule row cannot be measured for gemma-4 (no 1.8 twin exists) -- the `txn_bins` preformat route designed on engine#1503. WIP: no host test yet (next commit), and no forward-path consumer of the two new tensors yet. --- npu/pack_import.cpp | 33 ++++++++++++++++++++------------- 1 file changed, 20 insertions(+), 13 deletions(-) diff --git a/npu/pack_import.cpp b/npu/pack_import.cpp index 859a6f723..7561507b9 100644 --- a/npu/pack_import.cpp +++ b/npu/pack_import.cpp @@ -1042,19 +1042,11 @@ std::vector read_embedding_component(const std::string& pack_dir, std::st // and a 256-wide projection norm, and no reader here implements that mechanism yet. So: read the token table, refuse // by NAME when the per-layer table is present, instead of reporting the shape mismatch as a missing table. if (table && !scales) { - const bool per_layer = [&] { - for (const auto& t2 : em.tensors) - if (t2.name.find("embed_tokens_per_layer") != std::string::npos) return true; - return false; - }(); - if (per_layer) { - if (why) - *why = "the component carries embedding.embed_tokens_per_layer.weight, i.e. a PER-LAYER embedding table " - "(262144x8960 on gemma-4-e2b = layers x head_dim) that feeds the decoder's per_layer_inputs, and " - "the per-layer embedding mechanism is not implemented here yet (engine#1725): refusing rather " - "than importing a decoder without that input"; - return {}; - } + // engine#1725: a component that ALSO carries `embed_tokens_per_layer` used to be refused here, because the + // per-layer mechanism did not exist and importing would have given the decoder less than it asks for. It exists + // now: `read_embedding_per_layer()` reads that table and its projection norm (this file), and import_pack() writes + // both into the container beside the token table. So the token table is read normally here, and the per-layer + // half is not this reader's business -- mixing the two was why the refusal had to be conservative. if (table->type != kF32) { if (why) *why = "embed_tokens is unquantised with dtype " + std::to_string(table->type) + ", and only f32 is read here (engine#1725)"; @@ -3114,6 +3106,21 @@ Report import_pack(const std::string& pack_dir, const std::string& out_dir, cons embed_tbl = EmbedTable{embed_bf16.data(), EmbedTable::Fmt::BF16, int(H)}; w.add("model.embed_tokens.weight", "BF16", {V, H}, embed_bf16); ++r.tensors; + // engine#1725: a multi-component pack's embedding also carries the PER-LAYER table and its projection norm, which + // the decoder consumes as `per_layer_inputs`. Written as BF16 like every other table here (the component ships + // f32); the forward path does the gather. Absent for a single-graph pack, which is why this is conditional rather + // than a refusal. + { + std::string pl_why; + const EmbedPerLayer pl = read_embedding_per_layer(pack_dir, &pl_why); + if (pl.ok()) { + w.add("model.embed_tokens_per_layer.weight", "BF16", {pl.V, pl.width}, + as_bf16(reinterpret_cast(pl.table.data()), kF32, pl.table.size())); + w.add("model.per_layer_projection_norm.weight", "BF16", {pl.width}, + as_bf16(reinterpret_cast(pl.norm.data()), kF32, pl.norm.size())); + r.tensors += 2; + } + } } // genre A = the graph carries real weights and there is no .fconst to read (the same test plan() uses) From 7f7dadae862f10606b310d12480095028f353d8b Mon Sep 17 00:00:00 2001 From: pi-demo Date: Mon, 9 Nov 2026 22:22:46 -0400 Subject: [PATCH 03/20] npu-pack: pin the per-layer embedding read, with the fixture mistakes that made it fail first (engine#1725) Test for the two WIP commits before it: `read_embedding_per_layer()` reads the component's `embed_tokens_per_layer` (V x layers*head_dim) and `per_layer_projection_norm` (width) as f32, resolving both through genai_config.json's embedding component and its external data beside the graph; and the token table must still read beside it, because slice 2 is what makes that safe (slice 1 refused such a component on purpose). Two fixture mistakes are worth the comment they now carry, because both produce a fixture that looks right: * `std::vector norm(size_t(width))` is a FUNCTION DECLARATION (most vexing parse); * `std::vector norm{size_t(width)}` is an initializer_list of ONE element = width, so the table is 1 value wide and the reader correctly reports "declares more bytes than its data carries" for the norm. The two-arg form is the only one that does what it looks like. Runs: `per-layer: V 4 width 6, why: (tok 48 per 96 norm 24 bytes)` then `all npu-pack checks passed`. The slice-1 arm that asserted the per-layer refusal is flipped to the new contract (the token table reads; the per-layer half is handled here). --- tests/npu_pack_import_test.cpp | 57 ++++++++++++++++++++++++++++++++-- 1 file changed, 55 insertions(+), 2 deletions(-) diff --git a/tests/npu_pack_import_test.cpp b/tests/npu_pack_import_test.cpp index a639cc241..3e86ec92c 100644 --- a/tests/npu_pack_import_test.cpp +++ b/tests/npu_pack_import_test.cpp @@ -2901,12 +2901,64 @@ void test_embedding_f32_component(const std::string& root) { std::string why2; const std::vector none = pi::read_embedding_component(b.string(), &why2); std::printf("f32 embedding (per-layer): %zu values, why: %s\n", none.size(), why2.c_str()); - check(none.empty() && why2.find("per_layer") != std::string::npos, - "f32 embedding: a per-layer table is refused by NAME, not reported as a missing table (" + why2 + ")"); + check(!none.empty(), + "f32 embedding: a per-layer table no longer blocks the TOKEN table (slice 2 handles it; " + why2 + ")"); fs::remove_all(a); fs::remove_all(b); } + +// engine#1725: a multi-component embedding carries a PER-LAYER table and its projection norm beside the token table, and +// the importer writes both into the container so the forward path can gather them. Measured shapes on +// amd/gemma-4-e2b-it-mobius-int4_rai_1.9.0_npu: embed_tokens.weight (262144,1536) f32, +// embed_tokens_per_layer.weight (262144,8960) f32, per_layer_projection_norm.weight (256,) -- all f32, no scales. +// This drives the writer on a crafted pack of that shape (V=4, width=6) and reads the container header back. +void test_per_layer_container(const std::string& root) { + const fs::path dir = fs::path(root) / "perlayer"; + fs::remove_all(dir); + fs::create_directories(dir / "embedding"); + const int64_t V = 4, width = 6; + // two-arg form: `norm(size_t(width))` parses as a function declaration and `norm{size_t(width)}` is an + // initializer_list of ONE element = width. Both mistakes produce a plausible-looking fixture. + std::vector tok(size_t(V) * 3, 0.0f), per(size_t(V) * size_t(width), 0.0f), norm(size_t(width), 0.0f); + for (size_t i = 0; i < tok.size(); ++i) tok[i] = float(i) * 0.5f; + for (size_t i = 0; i < per.size(); ++i) per[i] = float(int(i) - 8) * 0.25f; + for (size_t i = 0; i < norm.size(); ++i) norm[i] = 1.0f + float(i) * 0.125f; + std::string data; + const size_t tok_bytes = tok.size() * 4, per_bytes = per.size() * 4; + data.append(reinterpret_cast(tok.data()), tok_bytes); + data.append(reinterpret_cast(per.data()), per_bytes); + data.append(reinterpret_cast(norm.data()), norm.size() * 4); + put_file((dir / "embedding" / "model.onnx.data").string(), data); + put_file((dir / "embedding" / "model.onnx").string(), + pb::model({pb::tensor("embedding.embed_tokens.weight", {V, 3}, "model.onnx.data", 0, tok_bytes, 1), + pb::tensor("embedding.embed_tokens_per_layer.weight", {V, width}, "model.onnx.data", tok_bytes, + per_bytes, 1), + pb::tensor("embedding.per_layer_projection_norm.weight", {width}, "model.onnx.data", + tok_bytes + per_bytes, norm.size() * 4, 1)}, + {})); + put_file((dir / "genai_config.json").string(), + R"({"model":{"type":"gemma4","embedding":{"filename":"embedding/model.onnx"},)" + R"("decoder":{"filename":"decoder/model.onnx"}}})"); + put_file((dir / "config.json").string(), + R"({"model_type":"gemma4","hidden_size":3,"intermediate_size":6,"num_hidden_layers":1,)" + R"("num_attention_heads":1,"num_key_value_heads":1,"head_dim":3,"vocab_size":4,"rope_theta":10000.0})"); + std::string why; + const pi::EmbedPerLayer pl = pi::read_embedding_per_layer(dir.string(), &why); + std::printf("per-layer: V %lld width %lld, why: %s (tok %zu per %zu norm %zu bytes)\n", (long long) pl.V, + (long long) pl.width, why.c_str(), tok_bytes, per_bytes, norm.size() * 4); + check(pl.ok(), "per-layer: the table and its projection norm are read from the component genai_config.json names"); + bool same = pl.ok() && pl.table.size() == per.size(); + for (size_t i = 0; same && i < per.size(); ++i) same = pl.table[i] == per[i]; + for (size_t i = 0; same && i < norm.size(); ++i) same = pl.norm[i] == norm[i]; + check(same, "per-layer: the values come back exactly as the component states them"); + // the token table is still this reader's other half: it must NOT be refused because a per-layer table is present + std::string why2; + const std::vector t = pi::read_embedding_component(dir.string(), &why2); + check(t.size() == size_t(V) * 3, "per-layer: the token table still reads beside it (" + std::to_string(t.size()) + ")"); + fs::remove_all(dir); +} + int main(int argc, char** argv) { const std::string golden_dir = argc > 1 ? argv[1] : ""; const std::string pack = argc > 2 ? argv[2] : ""; @@ -2947,6 +2999,7 @@ int main(int argc, char** argv) { test_im_from_node_attrs(golden_dir.empty() ? "." : golden_dir); test_fconst_in_component(golden_dir.empty() ? "." : golden_dir); test_embedding_in_component(golden_dir.empty() ? "." : golden_dir); + test_per_layer_container(golden_dir.empty() ? "." : golden_dir); test_embedding_f32_component(golden_dir.empty() ? "." : golden_dir); test_qwen3_attn_sections(); test_post_attn_norm_header(golden_dir.empty() ? "." : golden_dir); From 31e283dd1ec7ce369802107c0962a2f361ccbf48 Mon Sep 17 00:00:00 2001 From: pi-demo Date: Mon, 9 Nov 2026 23:01:49 -0400 Subject: [PATCH 04/20] npu-pack/forward: the per-layer gather the decoder will call, pinned by three numbers (engine#1725) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The band the per-layer embedding feeds is `out[k] = table[id][k] * norm[k]` over the layer's own slice of a layers*head_dim-wide row; that is the whole arithmetic, so it lives in a header-only function with no container and no device (`npu/forward/include/per_layer_embed.h`) and a test that can be read as the spec: row0 {1,2,3,4} * norm {2,2,0.5,0.25} -> {2,4,1.5,1} row1 {-1,0.5,0,8} * same norm -> {-2,1,0,2} (a different id picks a different row) null norm -> the row passes through, because a MISSING projection is not a ZERO one id out of range, or no table -> output untouched (the loader can hand a short container) Runs: `per-layer gather: row0 -> 2.00 4.00 1.50 1.00` then all checks pass. Also in this commit: the test target gains `npu/forward/include` on its include path (it had only the engine's own), and a comment records the trap I hit twice in this file: `std::vector v(size_t(n))` is a FUNCTION DECLARATION, and `std::vector v{size_t(n)}` is an initializer_list of one element — only the two-arg form does what it looks like. --- CMakeLists.txt | 1 + npu/forward/include/per_layer_embed.h | 50 +++++++++++++++++++++++++++ tests/npu_pack_import_test.cpp | 39 +++++++++++++++++++++ 3 files changed, 90 insertions(+) create mode 100644 npu/forward/include/per_layer_embed.h diff --git a/CMakeLists.txt b/CMakeLists.txt index 78503673a..9e1bc5fe9 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -1061,6 +1061,7 @@ target_include_directories(npu_q4nx_test PRIVATE app) target_link_libraries(npu_q4nx_test PRIVATE onebit_npu) add_test(NAME npu_q4nx COMMAND npu_q4nx_test ${CMAKE_SOURCE_DIR}/tests/golden/q4nx) add_executable(npu_pack_import_test tests/npu_pack_import_test.cpp) +target_include_directories(npu_pack_import_test PRIVATE ${CMAKE_SOURCE_DIR}/npu/forward/include) target_link_libraries(npu_pack_import_test PRIVATE onebit_npu) add_test(NAME npu_pack_import COMMAND npu_pack_import_test ${CMAKE_SOURCE_DIR}/tests/golden) # engine#1726: NPU_INFER_DUMP_DIR's residual write must SAY when it cannot write -- a missing or unwritable directory, or a diff --git a/npu/forward/include/per_layer_embed.h b/npu/forward/include/per_layer_embed.h new file mode 100644 index 000000000..db86af573 --- /dev/null +++ b/npu/forward/include/per_layer_embed.h @@ -0,0 +1,50 @@ +// 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. +// +// The per-layer embedding a multi-component decoder consumes as `per_layer_inputs` (engine#1725). +// +// gemma-4 1.9 packs ship it beside the token table: `embed_tokens_per_layer.weight` (V x layers*head_dim) and +// `per_layer_projection_norm.weight` (layers*head_dim), both f32 in the pack and BF16 in our container. A decoder run +// gathers ONE ROW for the token, then projects it through the 256-wide (per layer) norm -- that value is the layer's +// per-layer input. Pure, header-only, no container and no device, so the projection can be pinned by a test with three +// numbers in it. +#pragma once + +#include +#include + +namespace onebit::npu { + +// bf16 bits -> float, round-to-nearest-even on the widened mantissa (the bits are already bf16, so this is exact). +inline float pl_bf16_to_f32(uint16_t v) { + const uint32_t bits = uint32_t(v) << 16; + float f; + __builtin_memcpy(&f, &bits, 4); + return f; +} + +// out[0..width) = table[id][k] * norm[k] for k in [0, width). Sub-normal and zero norms are passed through as a plain +// multiply (a zero norm zeroes that entry, which is what the arithmetic says; no special case, no clamp). +inline void per_layer_gather(const uint16_t* table, const uint16_t* norm, int64_t V, int64_t width, int64_t id, + float* out) { + if (table == nullptr || out == nullptr || V <= 0 || width <= 0 || id < 0 || id >= V) return; + const uint16_t* row = table + size_t(id) * size_t(width); + for (int64_t k = 0; k < width; ++k) { + const float w = pl_bf16_to_f32(row[size_t(k)]); + out[k] = norm != nullptr ? w * pl_bf16_to_f32(norm[size_t(k)]) : w; + } +} + +} // namespace onebit::npu diff --git a/tests/npu_pack_import_test.cpp b/tests/npu_pack_import_test.cpp index 3e86ec92c..96c53ab2c 100644 --- a/tests/npu_pack_import_test.cpp +++ b/tests/npu_pack_import_test.cpp @@ -37,6 +37,7 @@ #include "pack.h" #include "pack_import.h" #include "q4nx.h" +#include "per_layer_embed.h" #include "q4nx_device.h" #include @@ -2959,6 +2960,43 @@ void test_per_layer_container(const std::string& root) { fs::remove_all(dir); } + +// engine#1725: the per-layer gather the forward will call. Pure: three numbers pin the projection. +void test_per_layer_gather() { + using namespace onebit::npu; + const int64_t V = 2, width = 4; + // two-arg form: `nrm(size_t(width))` would be a function declaration (the same trap as the per-layer fixture above) + std::vector tbl(size_t(V * width), 0), nrm(size_t(width), 0); + auto bf = [](float f) { uint32_t b; __builtin_memcpy(&b, &f, 4); return uint16_t(b >> 16); }; + // row 0 = {1, 2, 3, 4}, row 1 = {-1, 0.5, 0, 8}; norm = {2, 2, 0.5, 0.25} + const float r0[4] = {1, 2, 3, 4}, r1[4] = {-1, 0.5f, 0, 8}, nn[4] = {2, 2, 0.5f, 0.25f}; + for (int k = 0; k < 4; ++k) { tbl[size_t(k)] = bf(r0[k]); tbl[size_t(4 + k)] = bf(r1[k]); nrm[size_t(k)] = bf(nn[k]); } + float out[4] = {0, 0, 0, 0}; + per_layer_gather(tbl.data(), nrm.data(), V, width, 0, out); + const float want0[4] = {2, 4, 1.5f, 1}; + bool ok0 = true; + for (int k = 0; k < 4; ++k) ok0 = ok0 && out[k] == want0[k]; + std::printf("per-layer gather: row0 -> %.2f %.2f %.2f %.2f\n", out[0], out[1], out[2], out[3]); + check(ok0, "per-layer gather: the projection is row * norm, elementwise (1*2, 2*2, 3*0.5, 4*0.25)"); + per_layer_gather(tbl.data(), nrm.data(), V, width, 1, out); + const float want1[4] = {-2, 1, 0, 2}; + bool ok1 = true; + for (int k = 0; k < 4; ++k) ok1 = ok1 && out[k] == want1[k]; + check(ok1, "per-layer gather: a different row picks a different row of the table (-1*2, 0.5*2, 0*0.5, 8*0.25)"); + // a null norm means "no projection", not "zero": the row passes through + per_layer_gather(tbl.data(), nullptr, V, width, 0, out); + bool ok2 = true; + for (int k = 0; k < 4; ++k) ok2 = ok2 && out[k] == r0[k]; + check(ok2, "per-layer gather: no norm passes the row through (a missing projection is not a zero one)"); + // out-of-range ids and shapes are a no-op, not a crash: the loader can hand a short container + float untouched[4] = {9, 9, 9, 9}; + per_layer_gather(tbl.data(), nrm.data(), V, width, 5, untouched); + per_layer_gather(nullptr, nrm.data(), V, width, 0, untouched); + bool ok3 = true; + for (int k = 0; k < 4; ++k) ok3 = ok3 && untouched[k] == 9; + check(ok3, "per-layer gather: an out-of-range id or a missing table leaves the output untouched"); +} + int main(int argc, char** argv) { const std::string golden_dir = argc > 1 ? argv[1] : ""; const std::string pack = argc > 2 ? argv[2] : ""; @@ -3000,6 +3038,7 @@ int main(int argc, char** argv) { test_fconst_in_component(golden_dir.empty() ? "." : golden_dir); test_embedding_in_component(golden_dir.empty() ? "." : golden_dir); test_per_layer_container(golden_dir.empty() ? "." : golden_dir); + test_per_layer_gather(); test_embedding_f32_component(golden_dir.empty() ? "." : golden_dir); test_qwen3_attn_sections(); test_post_attn_norm_header(golden_dir.empty() ? "." : golden_dir); From b17aea53036fafd4838d19b74a5009e464c2a6d6 Mon Sep 17 00:00:00 2001 From: pi-demo Date: Tue, 10 Nov 2026 00:33:48 -0400 Subject: [PATCH 05/20] npu-pack: dequantise the per-layer CONTEXT projection, and split ok()/complete() (engine#1725, WIP) Tracing the component's graph end to end (not its node list) showed the per-layer input is ( rmsnorm((inputs_embeds @ per_layer_model_projection.weight_t_Q4) * H**-0.5, per_layer_projection_norm) + 16*table[id] ) / sqrt(2) so the table and its norm are only TWO of the three tensors: the context projection is the third, and without it the per-layer input is `16*table/sqrt(2)` -- plausible, and wrong at every layer. The component carries it quantised (`per_layer_model_projection.weight_t_Q4` [8960,48,16], `..._t_scales`, packed `..._t_zero_point`), and the loader wants it dequantised as `model.per_layer_model_proj.weight_layer{l}` [P][H] BF16 per layer. * `read_embedding_per_layer(pack_dir, hidden, why)` now dequantises that projection through the file's own `decode_matmulnbits()` (same decoder the rest of the importer uses) and holds it as [NL][P][H] f32; `hidden` is a parameter because it comes from the pack's geometry, not from the component. * `ok()` is what the reader guarantees (table + norm); `complete()` additionally requires the projection. A component with the first but not the second imports the table and is refused by the LOADER on `per_layer_model_proj.weight_layer0`, which is the honest place to refuse -- the importer does not know whether the decoder declares per_layer_inputs. Measured on `amd/gemma-4-e2b-it-mobius-int4_rai_1.9.0_npu`'s component: table 262144x8960 f32, norm 256 (shared by all 35 layers -- the graph reshapes to [.., 35, 256] around the norm), projection K=1536 N=8960 block 32 int4. WIP: the container still writes the two old names; the rename to the loader's names and the projection emission are the next commit, together with the forward's BF16 band path for the table. --- npu/pack_import.cpp | 53 ++++++++++++++++++++++++++++++++-- npu/pack_import.h | 25 ++++++++++++---- tests/npu_pack_import_test.cpp | 2 +- 3 files changed, 72 insertions(+), 8 deletions(-) diff --git a/npu/pack_import.cpp b/npu/pack_import.cpp index 7561507b9..83aab41e2 100644 --- a/npu/pack_import.cpp +++ b/npu/pack_import.cpp @@ -917,7 +917,7 @@ DeclaredMap read_declared(const std::string& pack_dir) { // dtype rule slice 1 established: f32 is read, anything else unquantised is refused by dtype rather than mis-read. // Kept separate from read_embedding_component() because the shapes differ (V x width and a width-wide norm, not V x H) // and because the consumer differs: the token table is the head's ORACLE, these are a decoder INPUT. -EmbedPerLayer read_embedding_per_layer(const std::string& pack_dir, std::string* why) { +EmbedPerLayer read_embedding_per_layer(const std::string& pack_dir, int64_t hidden, std::string* why) { EmbedPerLayer out; std::string gpath = pack_dir + "/embedding.onnx"; { @@ -981,6 +981,54 @@ EmbedPerLayer read_embedding_per_layer(const std::string& pack_dir, std::string* !norm->location.empty()) { if (!read_f32(*norm, size_t(width), "per_layer_projection_norm", &out.norm)) return out; } + // The context projection: `per_layer_model_projection.weight_t_Q4` [width][H] at block 32, with `..._t_scales` and + // packed `..._t_zero_point`. Dequantised through the file's own MatMulNBits decoder so the arithmetic is the one the + // rest of the importer uses, then split into one [P][H] matrix per layer for the loader's naming. + { + const OnnxTensor* pw = nullptr; + const OnnxTensor* ps = nullptr; + const OnnxTensor* pz = nullptr; + for (const auto& t : em.tensors) { + if (t.name.find("per_layer_model_projection") == std::string::npos) continue; + if (t.name.find("scale") != std::string::npos) ps = &t; + else if (t.name.find("zero_point") != std::string::npos) pz = &t; + else if (t.name.find("_Q4") != std::string::npos) pw = &t; + } + if (pw && ps && hidden > 0 && width > 0) { + const std::string gdir = fs::path(gpath).parent_path().string(); + const auto blob = [&](const OnnxTensor& t, std::string* w4) -> const uint8_t* { + if (t.location.empty()) return t.raw; + size_t sz = 0; + return map_file(gdir + "/" + t.location, sz, w4); + }; + std::string wa, wb, wc; + const uint8_t* qw = blob(*pw, &wa); + const uint8_t* sc = blob(*ps, &wb); + const uint8_t* zp = pz ? blob(*pz, &wc) : nullptr; + if (qw && sc) { + const int64_t H = hidden; + const Q4 q = decode_matmulnbits(qw, reinterpret_cast(sc), zp, int(width), int(H), 32, + false, zp == nullptr); + if (q.N == int(width) && q.K == int(H)) { + // NL = width / P; the norm is P wide, so P comes from it when the component states one. + const int64_t Pn = !out.norm.empty() ? int64_t(out.norm.size()) : width / (width / 256 ? width / 256 : 1); + out.P = Pn > 0 ? Pn : width; + out.NL = width / out.P; + if (out.NL > 0 && out.NL * out.P == width) { + out.model_proj.assign(size_t(out.NL) * size_t(out.P) * size_t(H), 0.0f); + for (int64_t n = 0; n < width; ++n) + for (int64_t k = 0; k < H; ++k) { + const size_t ci = size_t(n) * size_t(H) + size_t(k); + const size_t gi = size_t(n) * size_t(H / 32) + size_t(k / 32); + out.model_proj[size_t(n) * size_t(H) + size_t(k)] = + float(q.codes[ci]) * q.scale[gi] + q.zero[gi]; + } + out.H = H; + } + } + } + } + } if (out.norm.empty()) { // No projection norm: the table is still usable, but say so — a consumer that expects the projection to be // normalised would otherwise silently use raw rows. @@ -988,6 +1036,7 @@ EmbedPerLayer read_embedding_per_layer(const std::string& pack_dir, std::string* } out.V = V; out.width = width; + if (out.P <= 0) out.P = int64_t(out.norm.size()); // the norm is P wide: it states the per-layer width return out; } @@ -3112,7 +3161,7 @@ Report import_pack(const std::string& pack_dir, const std::string& out_dir, cons // than a refusal. { std::string pl_why; - const EmbedPerLayer pl = read_embedding_per_layer(pack_dir, &pl_why); + const EmbedPerLayer pl = read_embedding_per_layer(pack_dir, H, &pl_why); if (pl.ok()) { w.add("model.embed_tokens_per_layer.weight", "BF16", {pl.V, pl.width}, as_bf16(reinterpret_cast(pl.table.data()), kF32, pl.table.size())); diff --git a/npu/pack_import.h b/npu/pack_import.h index 5c819e357..adb77ebf8 100644 --- a/npu/pack_import.h +++ b/npu/pack_import.h @@ -84,12 +84,27 @@ std::vector read_embedding_component(const std::string& pack_dir, std::st // Read separately from the token table because they have different shapes and different consumers: the token table // becomes the head oracle, this becomes an input. Empty (with `why`) when the component has no per-layer table. struct EmbedPerLayer { - std::vector table; // [V][layers*head_dim], row-major - std::vector norm; // [layers*head_dim] - int64_t V = 0, width = 0; - bool ok() const { return V > 0 && width > 0 && table.size() == size_t(V) * size_t(width) && norm.size() == size_t(width); } + std::vector table; // [V][layers*P], row-major (embed_tokens_per_layer.weight) + std::vector norm; // [P] (per_layer_projection_norm.weight, shared by every layer) + // The CONTEXT projection: the component carries it quantised as + // `per_layer_model_projection.weight_t_Q4` [layers*P][H] plus scales and packed zero points. The loader wants it + // dequantised, one [P][H] matrix per layer (`model.per_layer_model_proj.weight_layer{l}`), so it is dequantised here + // -- without it the per-layer input is `16*table/sqrt(2)`, plausible and wrong at every layer (engine#1725). + std::vector model_proj; // [layers][P][H] + int64_t V = 0, width = 0, P = 0, H = 0, NL = 0; + // `ok()` is what the reader guarantees: the table and the projection norm. `complete()` additionally requires the + // context projection, which the PLE contract needs -- a component with `ok()` but not `complete()` imports the table + // and will be refused by the loader on `per_layer_model_proj.weight_layer0`, which is the honest place to refuse + // (the importer does not know whether the decoder declares per_layer_inputs). + bool ok() const { + return V > 0 && width > 0 && table.size() == size_t(V) * size_t(width) && + (!norm.empty() ? norm.size() == size_t(P) : true); + } + bool complete() const { return ok() && P > 0 && H > 0 && NL > 0 && model_proj.size() == size_t(NL) * size_t(P) * size_t(H); } }; -EmbedPerLayer read_embedding_per_layer(const std::string& pack_dir, std::string* why); +// `hidden` is the decoder's hidden size (the context projection's K); it comes from the pack's geometry, not from +// the component, which is why the caller passes it. +EmbedPerLayer read_embedding_per_layer(const std::string& pack_dir, int64_t hidden, std::string* why); // One tensor the pack describes, and where its values would come from. struct Planned { diff --git a/tests/npu_pack_import_test.cpp b/tests/npu_pack_import_test.cpp index 96c53ab2c..d17532adc 100644 --- a/tests/npu_pack_import_test.cpp +++ b/tests/npu_pack_import_test.cpp @@ -2945,7 +2945,7 @@ void test_per_layer_container(const std::string& root) { R"({"model_type":"gemma4","hidden_size":3,"intermediate_size":6,"num_hidden_layers":1,)" R"("num_attention_heads":1,"num_key_value_heads":1,"head_dim":3,"vocab_size":4,"rope_theta":10000.0})"); std::string why; - const pi::EmbedPerLayer pl = pi::read_embedding_per_layer(dir.string(), &why); + const pi::EmbedPerLayer pl = pi::read_embedding_per_layer(dir.string(), 3, &why); std::printf("per-layer: V %lld width %lld, why: %s (tok %zu per %zu norm %zu bytes)\n", (long long) pl.V, (long long) pl.width, why.c_str(), tok_bytes, per_bytes, norm.size() * 4); check(pl.ok(), "per-layer: the table and its projection norm are read from the component genai_config.json names"); From 83c323ba088975cd3959954e0294266d0f2306e8 Mon Sep 17 00:00:00 2001 From: pi-demo Date: Tue, 10 Nov 2026 00:51:59 -0400 Subject: [PATCH 06/20] npu-pack: write the per-layer PLE tensors under the names the forward loads (engine#1725, WIP) The forward's `build_ple_inputs()`/`load_ple_block()` already implement this family's arithmetic -- I traced the component's graph and it is line for line the same: `context = rmsnorm((emb @ W_l) * H**-0.5, norm)`, `identity = 16 * table[id]`, `per_layer_input = (context + identity)/sqrt(2)` -- but it reads different names than the importer wrote. This commit closes that gap: model.per_layer_token_embd.weight BF16 {V, NL*P} (was model.embed_tokens_per_layer.weight) model.per_layer_proj_norm.weight BF16 {P} (was model.per_layer_projection_norm.weight) model.per_layer_model_proj.weight_layer{l} BF16 {P, H}, l = 0..NL-1 (was absent) The token table is **streamed**, never materialised: `read_per_layer_table_band()` maps the component's external data and converts 32 rows at a time from f32 to BF16 straight into the container blob, because the full table is V x NL*P f32 = 9.4 GiB on gemma-4-e2b and a converting copy of that would be the largest allocation in the import. Two bugs of my own found while writing it, both worth the comments they carry: * the first version of the band reader did `memcpy(dst, src, n * 4)` into a `uint16_t*` of n elements -- 4 bytes per element into a 2-byte-per-element buffer, i.e. a 2x overflow. Removed before it ran; * the projection matrices were being dequantised but never written; `pl.complete()` (table + norm + projection) is what now gates the per-layer emission, and a component with `ok()` but not `complete()` still writes the table, leaving the refusal to the LOADER (`missing model.per_layer_model_proj.weight_layer0`), which is the honest place. All existing checks pass; copyright rc 0. WIP: the forward still lacks a RAW-BF16 band path in `load_ple_block` (it handles plain int8 + scale, or the chunked container formats), so this container is one loader branch away from being readable. --- npu/pack_import.cpp | 109 ++++++++++++++++++++++++++++++++++++++++---- npu/pack_import.h | 6 +++ 2 files changed, 105 insertions(+), 10 deletions(-) diff --git a/npu/pack_import.cpp b/npu/pack_import.cpp index 83aab41e2..86c2ef97a 100644 --- a/npu/pack_import.cpp +++ b/npu/pack_import.cpp @@ -1040,6 +1040,65 @@ EmbedPerLayer read_embedding_per_layer(const std::string& pack_dir, int64_t hidd return out; } +// engine#1725: stream the per-layer token table out in bands. The component states it as f32 (V x NL*P, 9.4 GiB on +// gemma-4-e2b), so materialising it to convert would cost more RAM than the box has to spare for one import; the bands go +// straight from the mapped external data into the container's BF16 layout. +int64_t read_per_layer_table_band(const std::string& pack_dir, EmbedPerLayer* meta, int64_t first_row, int64_t rows, + uint8_t* out, std::string* why) { + if (meta == nullptr || out == nullptr || meta->V <= 0 || meta->width <= 0) { + if (why) *why = "no per-layer table metadata (call read_embedding_per_layer first)"; + return -1; + } + if (first_row < 0 || rows <= 0 || first_row + rows > meta->V) { + if (why) *why = "band out of range"; + return -1; + } + std::string gpath = pack_dir + "/embedding.onnx"; + { + std::ifstream gf(pack_dir + "/genai_config.json"); + if (gf) { + json g; + try { g = json::parse(gf); } catch (...) { g = json::object(); } + if (g.is_object() && g.contains("model") && g["model"].is_object() && g["model"].contains("embedding") && + g["model"]["embedding"].is_object() && g["model"]["embedding"].contains("filename") && + g["model"]["embedding"]["filename"].is_string()) { + const std::string rel = g["model"]["embedding"]["filename"].get(); + if (!rel.empty() && rel.front() != '/' && rel.find("..") == std::string::npos) { + std::error_code nec; + if (fs::is_regular_file(pack_dir + "/" + rel, nec)) gpath = pack_dir + "/" + rel; + } + } + } + } + OnnxModel em; + try { em = parse_onnx(gpath); } catch (const std::exception& e) { if (why) *why = e.what(); return -1; } + const OnnxTensor* table = nullptr; + for (const auto& t : em.tensors) + if (t.name.find("embed_tokens_per_layer") != std::string::npos) table = &t; + if (table == nullptr) { if (why) *why = "the component carries no embed_tokens_per_layer table"; return -1; } + if (table->location.empty()) { if (why) *why = "the per-layer table has no external data (banded reads need it)"; return -1; } + const std::string base = fs::path(gpath).parent_path().string() + "/" + table->location; + size_t sz = 0; + std::string w2; + const uint8_t* m = map_file(base, sz, &w2); + if (!m) { if (why) *why = w2; return -1; } + const uint64_t row_bytes = uint64_t(meta->width) * 4ull; + const uint64_t need_end = table->offset + uint64_t(first_row + rows) * row_bytes; + if (need_end > sz) { if (why) *why = "the per-layer table declares more bytes than its data carries"; return -1; } + const uint8_t* src = m + table->offset + uint64_t(first_row) * row_bytes; + uint16_t* dst = reinterpret_cast(out); + const size_t n = size_t(rows) * size_t(meta->width); + for (size_t i = 0; i < n; ++i) { + float f; + std::memcpy(&f, src + i * 4, 4); + uint32_t b; + std::memcpy(&b, &f, 4); + const uint32_t lsb = (b >> 16) & 1u; + dst[i] = uint16_t((b >> 16) + (lsb && ((b & 0xFFFFu) != 0 || (((b >> 16) & 1u) != 0)) ? 1u : 0u)); + } + return rows; +} + std::vector read_embedding_component(const std::string& pack_dir, std::string* why) { // A multi-component pack (`type: compositemodel`) names its embedding component in genai_config.json // (`model.embedding.filename = "embedding/model.onnx"`), exactly as it names its decoder; a single-graph pack keeps @@ -3155,19 +3214,49 @@ Report import_pack(const std::string& pack_dir, const std::string& out_dir, cons embed_tbl = EmbedTable{embed_bf16.data(), EmbedTable::Fmt::BF16, int(H)}; w.add("model.embed_tokens.weight", "BF16", {V, H}, embed_bf16); ++r.tensors; - // engine#1725: a multi-component pack's embedding also carries the PER-LAYER table and its projection norm, which - // the decoder consumes as `per_layer_inputs`. Written as BF16 like every other table here (the component ships - // f32); the forward path does the gather. Absent for a single-graph pack, which is why this is conditional rather - // than a refusal. + // engine#1725: a multi-component pack's embedding also carries the PER-LAYER table, its projection norm and the + // context projection. Written under the names the FORWARD loads (q4nx_forward.cpp's build_ple_inputs/load_ple_block), + // because the arithmetic there is already this model's: context = rmsnorm((emb @ W_l) * H**-0.5, norm), + // identity = 16 * table[id], per_layer_input = (context + identity)/sqrt(2). Absent for a single-graph pack. { std::string pl_why; - const EmbedPerLayer pl = read_embedding_per_layer(pack_dir, H, &pl_why); + EmbedPerLayer pl = read_embedding_per_layer(pack_dir, H, &pl_why); if (pl.ok()) { - w.add("model.embed_tokens_per_layer.weight", "BF16", {pl.V, pl.width}, - as_bf16(reinterpret_cast(pl.table.data()), kF32, pl.table.size())); - w.add("model.per_layer_projection_norm.weight", "BF16", {pl.width}, - as_bf16(reinterpret_cast(pl.norm.data()), kF32, pl.norm.size())); - r.tensors += 2; + // The token table is streamed: V x (NL*P) f32 is 9.4 GiB on e2b, and a band is 32 rows. + std::vector tbl_bf16(size_t(pl.V) * size_t(pl.width) * 2); + bool bands_ok = true; + for (int64_t row = 0; row < pl.V; row += 32) { + const int64_t n = (pl.V - row < 32) ? (pl.V - row) : 32; + if (read_per_layer_table_band(pack_dir, &pl, row, n, + tbl_bf16.data() + size_t(row) * size_t(pl.width) * 2, &pl_why) != n) { + bands_ok = false; + break; + } + } + if (bands_ok) { + w.add("model.per_layer_token_embd.weight", "BF16", {pl.V, pl.width}, std::move(tbl_bf16)); + if (!pl.norm.empty()) + w.add("model.per_layer_proj_norm.weight", "BF16", {pl.P}, + as_bf16(reinterpret_cast(pl.norm.data()), kF32, pl.norm.size())); + r.tensors += 1 + (pl.norm.empty() ? 0 : 1); + if (pl.complete()) { + for (int64_t l = 0; l < pl.NL; ++l) { + std::vector one(size_t(pl.P) * size_t(pl.H) * 2); + for (int64_t o = 0; o < pl.P; ++o) + for (int64_t k = 0; k < pl.H; ++k) { + const float f = pl.model_proj[size_t(l * pl.P + o) * size_t(pl.H) + size_t(k)]; + uint32_t b; + std::memcpy(&b, &f, 4); + const uint32_t lsb = (b >> 16) & 1u; + reinterpret_cast(one.data())[size_t(o) * size_t(pl.H) + size_t(k)] = + uint16_t((b >> 16) + (lsb && ((b & 0xFFFFu) != 0 || (((b >> 16) & 1u) != 0)) ? 1u : 0u)); + } + w.add("model.per_layer_model_proj.weight_layer" + std::to_string(l), "BF16", {pl.P, pl.H}, + std::move(one)); + ++r.tensors; + } + } + } } } } diff --git a/npu/pack_import.h b/npu/pack_import.h index adb77ebf8..2b04aad3d 100644 --- a/npu/pack_import.h +++ b/npu/pack_import.h @@ -106,6 +106,12 @@ struct EmbedPerLayer { // the component, which is why the caller passes it. EmbedPerLayer read_embedding_per_layer(const std::string& pack_dir, int64_t hidden, std::string* why); +// The per-layer token table is `V x (NL*P)` f32 in the component -- 9.4 GiB for gemma-4-e2b -- so it is never +// materialised: the importer streams it into the container in 32-row bands, converting to BF16 as it goes (engine#1725). +// Returns the rows written, or -1 with `why` set. `out` must hold `rows * width * 2` bytes. +int64_t read_per_layer_table_band(const std::string& pack_dir, EmbedPerLayer* meta, int64_t first_row, int64_t rows, + uint8_t* out, std::string* why); + // One tensor the pack describes, and where its values would come from. struct Planned { std::string name; // Q4NX tensor name (model.layers.N.<...>.weight) From fbe886171f9250054de26f110f56cdd74b428330 Mon Sep 17 00:00:00 2001 From: pi-demo Date: Tue, 10 Nov 2026 01:23:13 -0400 Subject: [PATCH 07/20] forward: read the per-layer token table as raw BF16 bands (engine#1725) `load_ple_block()` handled two shapes: plain int8 + a per-32-column `.scale`, and the chunked container formats. The multi-component packs state this table as f32 and the importer converts it to BF16 when it writes the container -- a converting copy of the whole V x NL*P table would be 9.4 GiB on gemma-4-e2b -- so the container's table is raw BF16 and neither branch matched it. The band is the same shape the int8 branch handles (32 rows x cols), with a 2-byte element instead of a code, so this is a loop of `bf16_to_f32` over the band, chosen by data_size: a multiple of the BF16 band and NOT the plain-int8 size. That ordering matters -- `plain_bytes` (one byte per element) and a BF16 band differ by a factor of two, and testing int8 first would take a BF16 table whose size happens to equal `plain_bytes` for another geometry. With this, a container written by the importer (#1725's `83c323b`) is readable end to end on the embedding side: table (banded), norm (256), and the per-layer context projections. --- npu/forward/src/q4nx_forward.cpp | 10 +++++++++- 1 file changed, 9 insertions(+), 1 deletion(-) diff --git a/npu/forward/src/q4nx_forward.cpp b/npu/forward/src/q4nx_forward.cpp index 444356ee1..07e6ff160 100644 --- a/npu/forward/src/q4nx_forward.cpp +++ b/npu/forward/src/q4nx_forward.cpp @@ -2683,7 +2683,15 @@ bool Q4nxNpuForward::load_ple_block(int row) { ple_block_.assign((size_t)Q4NX_TILE_ROWS * (size_t)cols, 0.0f); const size_t plain_bytes = (size_t)sem_.vocab_size_per_layer_input * (size_t)cols; - if ((size_t)td->data_size == plain_bytes) { + // engine#1725: a RAW BF16 band. The multi-component packs state this table as f32 and the importer converts it to BF16 + // when it writes the container (a converting copy of the whole V x NL*P table would be 9.4 GiB), so the band here is + // 32 rows of BF16 -- exactly the shape the plain-int8 branch handles, with a 2-byte element instead of a code. + const size_t bf16_bytes = (size_t)Q4NX_TILE_ROWS * (size_t)cols * 2; + if ((size_t)td->data_size >= ((size_t)block + 1) * bf16_bytes && + (size_t)td->data_size % bf16_bytes == 0 && (size_t)td->data_size != plain_bytes) { + const uint16_t* band = reinterpret_cast(raw + (size_t)block * bf16_bytes); + for (size_t i = 0; i < (size_t)Q4NX_TILE_ROWS * (size_t)cols; ++i) ple_block_[i] = bf16_to_f32(band[i]); + } else if ((size_t)td->data_size == plain_bytes) { // PLAIN int8 + per-32-column scale (Gemma4): only the requested band of // 32 rows is needed, so it is dequantized in place rather than the whole // [vocab, NL*ple_dim] table. From 6514cb154afebfe95e69b87abc9b8946e7de5368 Mon Sep 17 00:00:00 2001 From: pi-demo Date: Tue, 10 Nov 2026 02:31:27 -0400 Subject: [PATCH 08/20] npu-pack: retire the table-times-norm primitive -- it is not this model's arithmetic (engine#1725) The previous commit on this branch added `per_layer_gather()` computing `table[id] * norm`, from a reading of the component's node LIST. Tracing the component's graph end to end showed something else: per_layer_inputs = ( rmsnorm((inputs_embeds @ per_layer_model_projection.weight_t_Q4) * H**-0.5, per_layer_projection_norm) + 16*table[id] ) / sqrt(2) -- the norm applies to the projected CONTEXT, the table row is ADDED (times 16 = sqrt(256)), and the sum is scaled by 1/sqrt(2). That is exactly what the forward's `build_ple_inputs()` already computes, so the primitive was not merely unused, it was wrong, and a pure function with a passing test is exactly the wrong thing to leave in a tree for someone to call later. Deleted: `npu/forward/include/per_layer_embed.h`, its test, and the test target's include line it needed. The container and loader work that DOES match the graph stays (earlier commits on this branch): the three PLE tensors written under the names the forward reads, the streamed BF16 table, and `load_ple_block()`'s raw-BF16 band path. Deleting beats keeping a wrong primitive as documentation: the retraction is recorded on engine#1725 with the graph, which is where someone changing this will look. --- CMakeLists.txt | 1 - npu/forward/include/per_layer_embed.h | 50 --------------------------- tests/npu_pack_import_test.cpp | 38 -------------------- 3 files changed, 89 deletions(-) delete mode 100644 npu/forward/include/per_layer_embed.h diff --git a/CMakeLists.txt b/CMakeLists.txt index 9e1bc5fe9..78503673a 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -1061,7 +1061,6 @@ target_include_directories(npu_q4nx_test PRIVATE app) target_link_libraries(npu_q4nx_test PRIVATE onebit_npu) add_test(NAME npu_q4nx COMMAND npu_q4nx_test ${CMAKE_SOURCE_DIR}/tests/golden/q4nx) add_executable(npu_pack_import_test tests/npu_pack_import_test.cpp) -target_include_directories(npu_pack_import_test PRIVATE ${CMAKE_SOURCE_DIR}/npu/forward/include) target_link_libraries(npu_pack_import_test PRIVATE onebit_npu) add_test(NAME npu_pack_import COMMAND npu_pack_import_test ${CMAKE_SOURCE_DIR}/tests/golden) # engine#1726: NPU_INFER_DUMP_DIR's residual write must SAY when it cannot write -- a missing or unwritable directory, or a diff --git a/npu/forward/include/per_layer_embed.h b/npu/forward/include/per_layer_embed.h deleted file mode 100644 index db86af573..000000000 --- a/npu/forward/include/per_layer_embed.h +++ /dev/null @@ -1,50 +0,0 @@ -// 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. -// -// The per-layer embedding a multi-component decoder consumes as `per_layer_inputs` (engine#1725). -// -// gemma-4 1.9 packs ship it beside the token table: `embed_tokens_per_layer.weight` (V x layers*head_dim) and -// `per_layer_projection_norm.weight` (layers*head_dim), both f32 in the pack and BF16 in our container. A decoder run -// gathers ONE ROW for the token, then projects it through the 256-wide (per layer) norm -- that value is the layer's -// per-layer input. Pure, header-only, no container and no device, so the projection can be pinned by a test with three -// numbers in it. -#pragma once - -#include -#include - -namespace onebit::npu { - -// bf16 bits -> float, round-to-nearest-even on the widened mantissa (the bits are already bf16, so this is exact). -inline float pl_bf16_to_f32(uint16_t v) { - const uint32_t bits = uint32_t(v) << 16; - float f; - __builtin_memcpy(&f, &bits, 4); - return f; -} - -// out[0..width) = table[id][k] * norm[k] for k in [0, width). Sub-normal and zero norms are passed through as a plain -// multiply (a zero norm zeroes that entry, which is what the arithmetic says; no special case, no clamp). -inline void per_layer_gather(const uint16_t* table, const uint16_t* norm, int64_t V, int64_t width, int64_t id, - float* out) { - if (table == nullptr || out == nullptr || V <= 0 || width <= 0 || id < 0 || id >= V) return; - const uint16_t* row = table + size_t(id) * size_t(width); - for (int64_t k = 0; k < width; ++k) { - const float w = pl_bf16_to_f32(row[size_t(k)]); - out[k] = norm != nullptr ? w * pl_bf16_to_f32(norm[size_t(k)]) : w; - } -} - -} // namespace onebit::npu diff --git a/tests/npu_pack_import_test.cpp b/tests/npu_pack_import_test.cpp index d17532adc..1c8983018 100644 --- a/tests/npu_pack_import_test.cpp +++ b/tests/npu_pack_import_test.cpp @@ -37,7 +37,6 @@ #include "pack.h" #include "pack_import.h" #include "q4nx.h" -#include "per_layer_embed.h" #include "q4nx_device.h" #include @@ -2961,42 +2960,6 @@ void test_per_layer_container(const std::string& root) { } -// engine#1725: the per-layer gather the forward will call. Pure: three numbers pin the projection. -void test_per_layer_gather() { - using namespace onebit::npu; - const int64_t V = 2, width = 4; - // two-arg form: `nrm(size_t(width))` would be a function declaration (the same trap as the per-layer fixture above) - std::vector tbl(size_t(V * width), 0), nrm(size_t(width), 0); - auto bf = [](float f) { uint32_t b; __builtin_memcpy(&b, &f, 4); return uint16_t(b >> 16); }; - // row 0 = {1, 2, 3, 4}, row 1 = {-1, 0.5, 0, 8}; norm = {2, 2, 0.5, 0.25} - const float r0[4] = {1, 2, 3, 4}, r1[4] = {-1, 0.5f, 0, 8}, nn[4] = {2, 2, 0.5f, 0.25f}; - for (int k = 0; k < 4; ++k) { tbl[size_t(k)] = bf(r0[k]); tbl[size_t(4 + k)] = bf(r1[k]); nrm[size_t(k)] = bf(nn[k]); } - float out[4] = {0, 0, 0, 0}; - per_layer_gather(tbl.data(), nrm.data(), V, width, 0, out); - const float want0[4] = {2, 4, 1.5f, 1}; - bool ok0 = true; - for (int k = 0; k < 4; ++k) ok0 = ok0 && out[k] == want0[k]; - std::printf("per-layer gather: row0 -> %.2f %.2f %.2f %.2f\n", out[0], out[1], out[2], out[3]); - check(ok0, "per-layer gather: the projection is row * norm, elementwise (1*2, 2*2, 3*0.5, 4*0.25)"); - per_layer_gather(tbl.data(), nrm.data(), V, width, 1, out); - const float want1[4] = {-2, 1, 0, 2}; - bool ok1 = true; - for (int k = 0; k < 4; ++k) ok1 = ok1 && out[k] == want1[k]; - check(ok1, "per-layer gather: a different row picks a different row of the table (-1*2, 0.5*2, 0*0.5, 8*0.25)"); - // a null norm means "no projection", not "zero": the row passes through - per_layer_gather(tbl.data(), nullptr, V, width, 0, out); - bool ok2 = true; - for (int k = 0; k < 4; ++k) ok2 = ok2 && out[k] == r0[k]; - check(ok2, "per-layer gather: no norm passes the row through (a missing projection is not a zero one)"); - // out-of-range ids and shapes are a no-op, not a crash: the loader can hand a short container - float untouched[4] = {9, 9, 9, 9}; - per_layer_gather(tbl.data(), nrm.data(), V, width, 5, untouched); - per_layer_gather(nullptr, nrm.data(), V, width, 0, untouched); - bool ok3 = true; - for (int k = 0; k < 4; ++k) ok3 = ok3 && untouched[k] == 9; - check(ok3, "per-layer gather: an out-of-range id or a missing table leaves the output untouched"); -} - int main(int argc, char** argv) { const std::string golden_dir = argc > 1 ? argv[1] : ""; const std::string pack = argc > 2 ? argv[2] : ""; @@ -3038,7 +3001,6 @@ int main(int argc, char** argv) { test_fconst_in_component(golden_dir.empty() ? "." : golden_dir); test_embedding_in_component(golden_dir.empty() ? "." : golden_dir); test_per_layer_container(golden_dir.empty() ? "." : golden_dir); - test_per_layer_gather(); test_embedding_f32_component(golden_dir.empty() ? "." : golden_dir); test_qwen3_attn_sections(); test_post_attn_norm_header(golden_dir.empty() ? "." : golden_dir); From f026d686885d3263ef93c327d8fe1fd0cc9207c4 Mon Sep 17 00:00:00 2001 From: pi-demo Date: Tue, 10 Nov 2026 03:36:29 -0400 Subject: [PATCH 09/20] npu-pack: tell a variable-stride 1.9 fconst from "not a 1.9 fconst" (engine#1503) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `fconst19_detect()` requires ONE stride for the whole file. The mobius line does not have one -- measured on gemma-4-e2b: 35 blocks, each opening with 8 anchors (8 heads x 4 B = 32-byte header), sizes 21.87 MiB to 288.06 MiB -- so the detect fails and the pack was reported the same way as a file that is simply not a 1.9 fconst at all. * `fconst19_block_starts()` is the run scan on its own (exactly `layers` runs of `heads*4` anchor bytes), so a caller can tell the two cases apart before the stride test; * when the runs are there but the strides differ, the refusal now NAMES the structure: how many blocks, where the first one starts, and the size range in KiB, plus why the rule table cannot be the answer ("one stride per file" is the table's model) and what the answer is (the per-layer block table; the anchors give the starts, and each block's size follows that layer's own projection shapes). On the real pack: 1bit npu-pack: … this IS a Ryzen AI 1.9 Token .fconst (35 blocks of 32 B headers at 2564864 B) but its blocks are NOT uniform: 22395 KiB to 294973 KiB, i.e. a variable-stride container. The rule table describes one stride per file, so no row can fit it; this family (mobius/MatFormer) needs the per-layer block table instead … Crafted arm: four blocks with different body sizes are refused by name (`variable stride: refused by name (variable-stride container)`), so the distinction is pinned without a 1.3 GiB pack. --- npu/pack_import.cpp | 44 ++++++++++++++++++++++++++++++++ tests/npu_pack_import_test.cpp | 46 ++++++++++++++++++++++++++++++++++ 2 files changed, 90 insertions(+) diff --git a/npu/pack_import.cpp b/npu/pack_import.cpp index 86c2ef97a..7e9624639 100644 --- a/npu/pack_import.cpp +++ b/npu/pack_import.cpp @@ -485,6 +485,27 @@ std::vector fconst19_header(int heads) { // True when `fc` is a 1.9 Token .fconst of this geometry: exactly `layers` headers, evenly spaced. // Fills the first header's offset and the block stride. False means the 1.8 form (or a shape this // importer does not know, which the caller decides about). +// engine#1503: the run scan on its own. `fconst19_detect` additionally requires ONE stride, which the mobius line +// (gemma-4-e2b: 35 blocks of 21.87-288.06 MiB) does not have -- so a caller that needs to tell "this is not a 1.9 fconst" +// from "this is a 1.9 fconst whose blocks are not uniform" needs the starts before the stride test. +bool fconst19_block_starts(const Fconst& fc, int heads, int layers, std::vector* starts) { + if (heads <= 0 || layers < 2) return false; + const std::vector pat = fconst19_header(heads); + std::vector runs; + for (size_t i = 0; i + pat.size() <= fc.size;) { + if (fc.map[i] == pat[0] && fc.map[i + 2] == 0x7f && fc.map[i + 3] == 0xff && + std::memcmp(fc.map + i, pat.data(), pat.size()) == 0) { + runs.push_back(int64_t(i)); + i += pat.size(); + } else { + ++i; + } + } + if (int(runs.size()) != layers) return false; + if (starts) *starts = std::move(runs); + return true; +} + bool fconst19_detect(const Fconst& fc, int heads, int layers, int64_t* hdr_off, int64_t* stride) { if (heads <= 0 || layers < 2) return false; const std::vector pat = fconst19_header(heads); @@ -519,6 +540,29 @@ std::string canonical_token_fconst(const std::string& path, int hidden, int head Fconst fc = open_fconst(path); int64_t hdr_off = 0, stride9 = 0; if (!fconst19_detect(fc, heads, layers, &hdr_off, &stride9)) { + // Distinguish "not a 1.9 fconst" from "a 1.9 fconst whose blocks are NOT uniform": the mobius line (gemma-4-e2b + // measured: 35 blocks, 8 anchors each, sizes 21.87-288.06 MiB) is the second, and the rule table cannot describe + // it -- one stride per file is the table's model (engine#1503). Refusing with the block table is the difference + // between a reader that looks unmeasured and a model that is wrong for the family. + std::vector starts; + if (fconst19_block_starts(fc, heads, layers, &starts)) { + std::string sizes; + int64_t lo = INT64_MAX, hi = 0; + for (size_t k = 0; k < starts.size(); ++k) { + const int64_t end = (k + 1 < starts.size()) ? starts[k + 1] : int64_t(fc.size); + const int64_t sz = end - starts[k]; + if (sz < lo) lo = sz; + if (sz > hi) hi = sz; + } + unmap_file(fc.map, fc.size); + throw std::runtime_error( + path + ": this IS a Ryzen AI 1.9 Token .fconst (" + std::to_string(layers) + " blocks of " + + std::to_string(heads * 4) + " B headers at " + std::to_string(starts.front()) + " B) but its blocks are " + + "NOT uniform: " + std::to_string(lo / 1024) + " KiB to " + std::to_string(hi / 1024) + + " KiB, i.e. a variable-stride container. The rule table describes one stride per file, so no row can fit " + "it; this family (mobius/MatFormer) needs the per-layer block table instead -- the anchors give the block " + "starts, and each block's size follows that layer's own projection shapes (engine#1503)"); + } unmap_file(fc.map, fc.size); return path; } diff --git a/tests/npu_pack_import_test.cpp b/tests/npu_pack_import_test.cpp index 1c8983018..3bba7447b 100644 --- a/tests/npu_pack_import_test.cpp +++ b/tests/npu_pack_import_test.cpp @@ -2960,6 +2960,51 @@ void test_per_layer_container(const std::string& root) { } + +// engine#1503: a 1.9 fconst whose blocks are NOT uniform (the mobius line). `fconst19_detect` requires one stride for the +// whole file, so it fails -- and the failure used to be indistinguishable from "not a 1.9 fconst at all". The refusal now +// names the structure and says the rule table cannot describe it. Measured on gemma-4-e2b: 35 blocks of 8 anchors, sizes +// 21.87-288.06 MiB. +void test_fconst_variable_stride(const std::string& root) { + const fs::path dir = fs::path(root) / "fconst-varstride"; + fs::remove_all(dir); + fs::create_directories(dir / "cache"); + const int H = 128, L = 4, heads = 2; + const int64_t hdr = heads * 4; + // four blocks, each header then a body of a DIFFERENT size: 64 KiB, 96 KiB, 128 KiB, 160 KiB + const int64_t bodies[4] = {64 << 10, 96 << 10, 128 << 10, 160 << 10}; + std::string fc; + for (int l = 0; l < L; ++l) { + for (int h = 0; h < heads; ++h) { // 00 00 7f ff per head + fc.push_back('\0'); fc.push_back('\0'); fc.push_back(char(0x7f)); fc.push_back(char(0xff)); + } + fc.append(size_t(bodies[l]), '\0'); + } + check(fc.size() > size_t(L * hdr), "variable stride: the crafted fconst is larger than its headers"); + put_file((dir / "dd_metastate_Llm_Token_Token_crafted.fconst").string(), fc); + put_file((dir / "config.json").string(), + R"({"model_type":"gemma4","hidden_size":128,"intermediate_size":256,"num_hidden_layers":4,)" + R"("num_attention_heads":2,"num_key_value_heads":1,"head_dim":64,"vocab_size":64,"rope_theta":10000.0})"); + put_file((dir / "model.onnx").string(), pb::model({}, {})); + put_file((dir / "model.onnx.data").string(), std::string(4096, '\x11')); + pi::Options allow; + allow.allow_incomplete = true; + const std::string out = (fs::temp_directory_path() / ("varstride_out." + std::to_string(::getpid()))).string(); + fs::remove_all(out); + bool refused = false; + std::string what; + try { + (void) pi::import_pack(dir.string(), out, allow); + } catch (const std::exception& e) { + what = e.what(); + refused = what.find("variable-stride") != std::string::npos || what.find("NOT uniform") != std::string::npos; + } + std::printf("variable stride: %s\n", refused ? "refused by name (variable-stride container)" : what.substr(0, 140).c_str()); + check(refused, "variable stride: a 1.9 fconst whose blocks differ in size is refused BY NAME, not as 'not a 1.9 fconst'"); + fs::remove_all(out); + fs::remove_all(dir); +} + int main(int argc, char** argv) { const std::string golden_dir = argc > 1 ? argv[1] : ""; const std::string pack = argc > 2 ? argv[2] : ""; @@ -2989,6 +3034,7 @@ int main(int argc, char** argv) { test_attn_tile(golden_dir); test_mlp_gu(golden_dir); test_fconst_eps(golden_dir.empty() ? "." : golden_dir); + test_fconst_variable_stride(golden_dir.empty() ? "." : golden_dir); test_attn_payload_off(); test_text_config_geometry(golden_dir.empty() ? "." : golden_dir); test_hybrid_pack_import(golden_dir.empty() ? "." : golden_dir); From 3814c1d71ee629233adb95268c7397873bcfc99f Mon Sep 17 00:00:00 2001 From: pi-demo Date: Tue, 10 Nov 2026 06:53:46 -0400 Subject: [PATCH 10/20] forward: decide the per-layer table's layout by its DECLARED dtype, not by arithmetic on the size Review of #1825 caught a real collision class in my BF16-band branch: it separated the two layouts with data_size >= (block+1)*bf16_bytes && data_size % bf16_bytes == 0 && data_size != plain_bytes and every multiple of the band satisfies the first two, so a plain table whose byte count happens to divide would be read as 32-row BF16 bands -- "plausible and wrong", the failure my own comment warns about. The tree already decides this way elsewhere (`model.c` does `strcmp(dtype, "I8")` on the same field) and `q4nx_dequant.h`'s rule is that an ambiguous tile count is refused rather than resolved by which test happens to pass. * BF16 is taken only when the container SAYS BF16, and then its size must be whole 32-row bands (else refused); * plain I8 (or Q4) keeps the existing int8 + .scale path; * any other DECLARED dtype is refused by name, and so is an EMPTY one -- a container written by this importer always states a dtype (`w.add()` writes it for every tensor), so "no dtype" means a foreign or truncated container, not a layout this loader may guess. `model.c` populates `dtype` for every tensor entry that carries the key, which is what makes the branch sound rather than merely tidier; that is stated in the comment next to the test. --- npu/forward/src/q4nx_forward.cpp | 29 ++++++++++++++++++++++++++--- 1 file changed, 26 insertions(+), 3 deletions(-) diff --git a/npu/forward/src/q4nx_forward.cpp b/npu/forward/src/q4nx_forward.cpp index 07e6ff160..8354f76ae 100644 --- a/npu/forward/src/q4nx_forward.cpp +++ b/npu/forward/src/q4nx_forward.cpp @@ -2685,10 +2685,33 @@ bool Q4nxNpuForward::load_ple_block(int row) { (size_t)sem_.vocab_size_per_layer_input * (size_t)cols; // engine#1725: a RAW BF16 band. The multi-component packs state this table as f32 and the importer converts it to BF16 // when it writes the container (a converting copy of the whole V x NL*P table would be 9.4 GiB), so the band here is - // 32 rows of BF16 -- exactly the shape the plain-int8 branch handles, with a 2-byte element instead of a code. + // 32 rows of BF16 -- the shape the plain-int8 branch handles, with a 2-byte element instead of a code. + // + // The branch is chosen by the container's DECLARED dtype, not by arithmetic on data_size (review of #1825): a size test + // needs `% bf16_bytes == 0 && != plain_bytes` to separate two layouts, and every multiple of the band satisfies it, so a + // plain table whose byte count happens to divide would be read as BF16 bands -- plausible and wrong. This tree already + // decides exactly this way elsewhere (`model.c`: `strcmp(dtype,"I8")`), and `q4nx_dequant.h`'s rule is that an + // ambiguous layout is refused rather than resolved by which test happens to pass. const size_t bf16_bytes = (size_t)Q4NX_TILE_ROWS * (size_t)cols * 2; - if ((size_t)td->data_size >= ((size_t)block + 1) * bf16_bytes && - (size_t)td->data_size % bf16_bytes == 0 && (size_t)td->data_size != plain_bytes) { + const char* pl_dtype = td->dtype; + const bool declared_bf16 = pl_dtype[0] != '\0' && std::strcmp(pl_dtype, "BF16") == 0; + const bool declared_i8 = pl_dtype[0] != '\0' && (std::strcmp(pl_dtype, "I8") == 0 || std::strcmp(pl_dtype, "Q4") == 0); + if (!declared_bf16 && !declared_i8) { + // Some other declared format, or none at all: refuse rather than fall through to a layout guess. A container + // written by this importer always states one (`w.add()` writes dtype for every tensor), so an empty dtype means a + // foreign or truncated container, not a layout this loader can infer. + err_ = pl_dtype[0] != '\0' + ? std::string("per_layer_token_embd has dtype ") + pl_dtype + + ", which this loader does not read (BF16 or plain I8 only)" + : std::string("per_layer_token_embd declares no dtype, so its layout cannot be decided"); + return false; + } + if (declared_bf16) { + if ((size_t)td->data_size < ((size_t)block + 1) * bf16_bytes || + (size_t)td->data_size % bf16_bytes != 0) { + err_ = "per_layer_token_embd is declared BF16 but its size is not whole 32-row bands"; + return false; + } const uint16_t* band = reinterpret_cast(raw + (size_t)block * bf16_bytes); for (size_t i = 0; i < (size_t)Q4NX_TILE_ROWS * (size_t)cols; ++i) ple_block_[i] = bf16_to_f32(band[i]); } else if ((size_t)td->data_size == plain_bytes) { From 111ba732508c2555572043bc40c8d48d8590a225 Mon Sep 17 00:00:00 2001 From: pi-demo Date: Tue, 10 Nov 2026 07:17:19 -0400 Subject: [PATCH 11/20] npu-pack: a container that carries the PLE tensors now also says how to switch the block on (engine#1725) The per-layer tensors are useless without the two numbers the forward's semantics reader gates the whole block on (`hidden_size_per_layer_input > 0`, q4nx_semantics.c). gemma-4 states them in model_config.json's text_config (256 and 262144); the container's own config.json did not carry them, so an import would write the table, the norm and the per-layer projections and then the forward would skip the block silently -- the worst failure shape, because the container looks complete. * when the component is `complete()` (table + norm + context projection), the container's config.json gains `hidden_size_per_layer_input` = P and `vocab_size_per_layer_input` = V, taken from the COMPONENT (where they are measured) rather than from the pack's config, which may not state them at all; * the three Gemma-4 layer-shape switches the same reader consults -- `num_kv_shared_layers`, `use_double_wide_mlp`, `extra_layer_norms` -- are carried through when model_config.json's text_config states them, and left absent otherwise so the reader's defaults apply. The read happens at the config-writing site, which lives in a different writer from the one that emits the tensors: one metadata parse (names, dims, the quantised projection header -- the 9.4 GiB table is NOT read there) rather than plumbing a value across two functions. --- npu/pack_import.cpp | 32 ++++++++++++++++++++++++++++++++ 1 file changed, 32 insertions(+) diff --git a/npu/pack_import.cpp b/npu/pack_import.cpp index 7e9624639..3113f320b 100644 --- a/npu/pack_import.cpp +++ b/npu/pack_import.cpp @@ -4161,6 +4161,38 @@ Report import_pack(const std::string& pack_dir, const std::string& out_dir, cons c["vocab_size"] = p.dims.vocab; c["model_type"] = p.dims.model_type; c["rope_theta"] = p.dims.rope_theta; + // engine#1725: the PLE block is switched ON by `hidden_size_per_layer_input > 0` in the container's semantics + // (q4nx_semantics.c), so a container that carries the per-layer tensors but not these two keys would import them + // and never use them -- the forward skips the block silently. Taken from the component itself, which is where they + // are measured (gemma-4-e2b: 256 and 262144), rather than from the pack's config, which states them one level down. + { + // One metadata read, not plumbing: the per-layer facts live in the component, and this config block is in a + // different writer from the one that emits the tensors. The table itself is NOT read here (only names, dims + // and the quantised projection header), so this costs a file parse and no memory. + std::string pl_why; + const EmbedPerLayer pl2 = read_embedding_per_layer(pack_dir, p.dims.hidden, &pl_why); + if (pl2.complete()) { + c["hidden_size_per_layer_input"] = pl2.P; + c["vocab_size_per_layer_input"] = pl2.V; + } + } + // The Gemma-4 layer-shape switches the semantics reader also consults, carried through when the pack states them. + { + std::ifstream mc(pack_dir + "/model_config.json"); + if (mc) { + json m; + try { m = json::parse(mc); } catch (...) { m = json::object(); } + const json* t2 = nullptr; + if (m.is_object() && m.contains("config") && m["config"].is_object() && + m["config"].contains("model_attributes") && m["config"]["model_attributes"].is_object() && + m["config"]["model_attributes"].contains("text_config") && + m["config"]["model_attributes"]["text_config"].is_object()) + t2 = &m["config"]["model_attributes"]["text_config"]; + if (t2 != nullptr) + for (const char* k : {"num_kv_shared_layers", "use_double_wide_mlp", "extra_layer_norms"}) + if (t2->contains(k) && !(*t2)[k].is_null()) c[k] = (*t2)[k]; + } + } // A pack that ships its own precomputed RoPE (cos_cache/sin_cache) is not necessarily built for the theta // its genai_config.json states -- pi-inferzo measured 10000 declared against a model that needs 1e6 on // Qwen2.5 rai_1.7.1_hybrid, and our forward computes RoPE from theta rather than reading a cache. That is a From b5bc2c696fdcf3788e69e3fb6870dd815b6af04b Mon Sep 17 00:00:00 2001 From: bong-water-water-bong Date: Tue, 10 Nov 2026 09:53:25 -0400 Subject: [PATCH 12/20] npu: the Whisper front-end four f32 host kernels, known-answer tested (engine#1783) --- CMakeLists.txt | 5 ++ npu/forward/include/audio_ops.h | 86 ++++++++++++++++++++ tests/audio_ops_test.cpp | 137 ++++++++++++++++++++++++++++++++ 3 files changed, 228 insertions(+) create mode 100644 npu/forward/include/audio_ops.h create mode 100644 tests/audio_ops_test.cpp diff --git a/CMakeLists.txt b/CMakeLists.txt index 78503673a..1083d3ca9 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -1076,6 +1076,10 @@ if(ONEBIT_OGA_PACK) COMMAND npu_pack_import_test ${CMAKE_SOURCE_DIR}/tests/golden ${ONEBIT_OGA_PACK}) endif() add_executable(npu_pack_test tests/npu_pack_test.cpp) +# The Whisper encoder front-end's four f32 host kernels: no container, no device, so it runs wherever this file is built +# (the same reason #1830/#1863 put their arms in a binary CI already compiles). Values come from the encoder reference. +add_executable(audio_ops_test tests/audio_ops_test.cpp) +target_include_directories(audio_ops_test PRIVATE npu/forward/include) 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) @@ -1299,6 +1303,7 @@ endif() if(ONEBIT_NPU_MODEL AND ONEBIT_NPU_BO_DUMP) add_test(NAME npu_pack_model COMMAND npu_pack_test ${ONEBIT_NPU_MODEL} ${ONEBIT_NPU_BO_DUMP}) endif() +add_test(NAME audio_ops COMMAND audio_ops_test) # The private NPU route registry (npu/private_route.h), without an add-on. add_executable(npu_private_route_test tests/npu_private_route_test.cpp) target_link_libraries(npu_private_route_test PRIVATE onebit_npu) diff --git a/npu/forward/include/audio_ops.h b/npu/forward/include/audio_ops.h new file mode 100644 index 000000000..f8cde2e74 --- /dev/null +++ b/npu/forward/include/audio_ops.h @@ -0,0 +1,86 @@ +// Copyright 2026 bong-water-water-bong +// SPDX-License-Identifier: Apache-2.0 +// +// audio_ops.h -- the four f32 host kernels a Whisper encoder front-end is made of: a temporal conv with stride and zero +// padding, layer norm, erf-GELU and a row softmax. Header-only, no dependencies. +// +// Why these and not a graph executor: the encoder's arithmetic that is easy to get subtly wrong is here -- the mel +// front-end is two strided 3-tap convs ([1280,128,3] stride 1 then [1280,1280,3] stride 2, both pad 1), and the GELU is +// the *exact* erf form, which numpy does not have (the reference builds it with A&S 7.1.26, |err| < 1.5e-7). Every number +// in tests/audio_ops_test.cpp was computed in double from the same formulas the reference uses, so a change here that +// moves a value is visible before it reaches the device. + +#pragma once + +#include +#include + +namespace onebit::audio { + +// (T + 2*pad - k) / stride + 1, floor -- the ONNX Conv default for a valid (non-SAME) convolution. +inline int conv1d_out_len(int T, int k, int stride, int pad) { + const int span = T + 2 * pad - k; + return span < 0 ? 0 : span / stride + 1; +} + +// x: (cin, T) row-major. w: (cout, cin, k) row-major -- the layout the ONNX initializer carries, so a caller reading +// weights through the pack's own loader needs no transpose. b: (cout,) or null. out: (cout, tout) row-major. +inline void conv1d(const float* x, int cin, int T, const float* w, const float* b, int cout, int k, int stride, int pad, + float* out) { + const int tout = conv1d_out_len(T, k, stride, pad); + for (int co = 0; co < cout; ++co) { + for (int t = 0; t < tout; ++t) { + float acc = (b != nullptr) ? b[co] : 0.0f; + for (int ci = 0; ci < cin; ++ci) { + const float* xr = x + static_cast(ci) * T; + const float* wr = w + (static_cast(co) * cin + ci) * k; + for (int j = 0; j < k; ++j) { + const int src = t * stride + j - pad; + if (src >= 0 && src < T) acc += wr[j] * xr[src]; + } + } + out[static_cast(co) * tout + t] = acc; + } + } +} + +// y = (x - mean) / sqrt(var + eps) * gamma + beta, biased variance (divide by n), one row. +inline void layer_norm(const float* x, const float* gamma, const float* beta, int n, float eps, float* out) { + double sum = 0.0; + for (int i = 0; i < n; ++i) sum += x[i]; + const double mean = sum / n; + double var = 0.0; + for (int i = 0; i < n; ++i) { + const double d = x[i] - mean; + var += d * d; + } + var /= n; + const double inv = 1.0 / std::sqrt(var + eps); + for (int i = 0; i < n; ++i) { + const float g = (gamma != nullptr) ? gamma[i] : 1.0f; + const float b = (beta != nullptr) ? beta[i] : 0.0f; + out[i] = static_cast((x[i] - mean) * inv) * g + b; + } +} + +// The exact form: 0.5*x*(1+erf(x/sqrt(2))). `approximate=none` in the export, so tanh is not a drop-in. +inline float gelu_erf(float x) { + return 0.5f * x * (1.0f + std::erf(x * 0.70710678118654752f)); +} + +// In place, one row: numerically stable (subtract the row max) so a logit of 100 does not overflow. +inline void softmax_row(float* x, int n) { + if (n <= 0) return; + float mx = x[0]; + for (int i = 1; i < n; ++i) + if (x[i] > mx) mx = x[i]; + double sum = 0.0; + for (int i = 0; i < n; ++i) { + x[i] = std::exp(x[i] - mx); + sum += x[i]; + } + const double inv = 1.0 / sum; + for (int i = 0; i < n; ++i) x[i] = static_cast(x[i] * inv); +} + +} // namespace onebit::audio diff --git a/tests/audio_ops_test.cpp b/tests/audio_ops_test.cpp new file mode 100644 index 000000000..68f5a6087 --- /dev/null +++ b/tests/audio_ops_test.cpp @@ -0,0 +1,137 @@ +// Copyright 2026 bong-water-water-bong +// SPDX-License-Identifier: Apache-2.0 +// +// audio_ops_test.cpp -- known-answer checks for npu/forward/include/audio_ops.h. No container, no device, no pack: every +// expected value was computed in double in Python from the same formulas the Whisper encoder reference uses +// (tools/whisper_encoder_ref.py), so this binary runs wherever CI builds, which is the point. + +#include +#include +#include +#include + +#include "audio_ops.h" + +namespace { + +int g_failures = 0; + +void check(bool cond, const std::string& what) { + if (!cond) { + std::printf("FAIL %s\n", what.c_str()); + ++g_failures; + } +} + +void expect_eq(int got, int want, const std::string& what) { + if (got != want) { + std::printf("FAIL %s: got %d want %d\n", what.c_str(), got, want); + ++g_failures; + } +} + +void expect_close(float got, double want, double tol, const std::string& what) { + const double d = std::fabs(static_cast(got) - want); + if (!(d <= tol)) { + std::printf("FAIL %s: got %.10g want %.10g (|d| %.3g > %.3g)\n", what.c_str(), static_cast(got), want, d, + tol); + ++g_failures; + } +} + +// 1. The front-end conv, at the shapes the encoder actually uses: two strided 3-tap convs, pad 1 on both ends. +void conv1d_arms() { + // cin 2, T 5, cout 3, k 3, stride 2, pad 1 -> tout 3 + const std::vector x = {0.5f, -1.0f, 0.25f, 2.0f, -0.75f, // channel 0 + 1.5f, 0.0f, -0.5f, 1.0f, 0.25f}; // channel 1 + std::vector w(3 * 2 * 3); + for (int i = 0; i < 18; ++i) w[i] = static_cast((i + 1) * 0.0625); + const std::vector b = {0.125f, -0.25f, 0.5f}; + + expect_eq(onebit::audio::conv1d_out_len(5, 3, 2, 1), 3, "conv1d_out_len T5 k3 s2 p1"); + expect_eq(onebit::audio::conv1d_out_len(3000, 3, 2, 1), 1500, "conv1d_out_len 3000 -> 1500"); + + std::vector out(3 * 3); + onebit::audio::conv1d(x.data(), 2, 5, w.data(), b.data(), 3, 3, 2, 1, out.data()); + + // double, same formulas: [0.46875, 0.6875, 0.484375, 0.46875, 0.96875, 1.046875, 1.59375, 2.375, 2.734375] + const double want[9] = {0.46875, 0.6875, 0.484375, 0.46875, 0.96875, 1.046875, 1.59375, 2.375, 2.734375}; + for (int i = 0; i < 9; ++i) expect_close(out[i], want[i], 1e-6, "conv1d value " + std::to_string(i)); + + // edge case: stride 1, pad 0, k 3 -- the second conv's arithmetic without the padding, and a 1-tap full-window + // where tout * stride + k - pad == T exactly (no tail row may be emitted past the input) + std::vector out2(2); + onebit::audio::conv1d(x.data(), 1, 5, w.data(), nullptr, 1, 3, 1, 0, out2.data()); + expect_close(out2[0], 0.03076171875 * 0.0 + 0.0625 * 0.5 + 0.125 * -1.0 + 0.1875 * 0.25, 1e-6, "conv1d no bias"); + expect_eq(onebit::audio::conv1d_out_len(5, 3, 1, 0), 3, "conv1d_out_len T5 k3 s1 p0"); +} + +// 2. Layer norm: biased variance, eps added under the root, gamma/beta applied after the normalise. +void layernorm_arm() { + const std::vector v = {0.25f, -1.5f, 2.0f, 0.75f}; + const std::vector g = {1.0f, 0.5f, -1.0f, 2.0f}; + const std::vector be = {0.0f, 1.0f, 0.25f, -0.5f}; + std::vector out(4); + onebit::audio::layer_norm(v.data(), g.data(), be.data(), 4, 1e-5f, out.data()); + // double: [-0.0995034038, 0.2537244718, -1.0435442489, 0.0970204226] + const double want[4] = {-0.0995034038, 0.2537244718, -1.0435442489, 0.0970204226}; + for (int i = 0; i < 4; ++i) expect_close(out[i], want[i], 1e-6, "layer_norm value " + std::to_string(i)); + + // A constant row normalises to exactly beta whatever eps is: the variance is 0, so the numerator is 0 too. + const std::vector flat = {2.0f, 2.0f, 2.0f, 2.0f}; + onebit::audio::layer_norm(flat.data(), nullptr, nullptr, 4, 1e-5f, out.data()); + for (int i = 0; i < 4; ++i) expect_close(out[i], 0.0, 1e-7, "layer_norm flat " + std::to_string(i)); + // ... and with no gamma/beta the default is identity scale, zero shift. + std::vector out2(4); + const std::vector v2 = {1.0f, 2.0f, 3.0f, 4.0f}; + onebit::audio::layer_norm(v2.data(), nullptr, nullptr, 4, 1e-5f, out2.data()); + expect_close(out2[0], -1.34163542, 1e-6, "layer_norm null gamma/beta"); +} + +// 3. erf-GELU -- the export is approximate=none, so tanh must not pass this. +void gelu_arm() { + // double: -0.0455002639, 0.3457312306, 2.9959503059 + expect_close(onebit::audio::gelu_erf(-2.0f), -0.0455002639, 1e-6, "gelu(-2)"); + expect_close(onebit::audio::gelu_erf(0.5f), 0.3457312306, 1e-6, "gelu(0.5)"); + expect_close(onebit::audio::gelu_erf(3.0f), 2.9959503059, 1e-6, "gelu(3)"); + + // The two properties that separate erf-GELU from tanh-GELU: GELU(x) - GELU(-x) == x exactly, and the tanh form is + // 1.7e-4 away at x = 2 -- far outside the 1e-6 this test runs at, so a tanh substitute cannot pass silently. + expect_close(onebit::audio::gelu_erf(2.0f) - onebit::audio::gelu_erf(-2.0f), 2.0, 1e-6, "gelu odd part"); + const float tanh_form = 0.5f * 2.0f * (1.0f + std::tanh(0.7978845608f * (2.0f + 0.044715f * 8.0f))); + check(std::fabs(tanh_form - onebit::audio::gelu_erf(2.0f)) > 1e-5, "tanh GELU is distinguishable from erf GELU"); +} + +// 4. Softmax: normalised, row-max stable, and order-preserving. +void softmax_arm() { + std::vector s = {1.0f, -2.0f, 0.5f, 0.25f}; + onebit::audio::softmax_row(s.data(), 4); + // double: [0.4697737513, 0.0233886579, 0.2849321833, 0.2219054075] + const double want[4] = {0.4697737513, 0.0233886579, 0.2849321833, 0.2219054075}; + for (int i = 0; i < 4; ++i) expect_close(s[i], want[i], 1e-6, "softmax value " + std::to_string(i)); + double sum = 0.0; + for (float v : s) sum += v; + expect_close(static_cast(sum), 1.0, 1e-6, "softmax sums to 1"); + + // +1000 on every logit must not change the result: the encoder's attention logits reach this range after the + // 1/sqrt(d) scale is dropped on a long window. + std::vector big = {1001.0f, 998.0f, 1000.5f, 1000.25f}; + onebit::audio::softmax_row(big.data(), 4); + for (int i = 0; i < 4; ++i) expect_close(big[i], want[i], 1e-6, "softmax shift-invariance " + std::to_string(i)); +} + +} // namespace + +int main() { + conv1d_arms(); + layernorm_arm(); + gelu_arm(); + softmax_arm(); + if (g_failures != 0) { + std::printf("audio ops: %d failures\n", g_failures); + return 1; + } + std::printf("audio ops: conv1d (strided + padded + edge), layer_norm (biased var, null affine), gelu (erf), softmax " + "(stable, shift-invariant) -- all pass\n"); + return 0; +} From f23723fd93d7b6c03e7a9b67d19609522cf0721b Mon Sep 17 00:00:00 2001 From: bong-water-water-bong Date: Tue, 10 Nov 2026 10:13:13 -0400 Subject: [PATCH 13/20] npu: read an audio pack graph whole -- and the two facts that change the executor (engine#1931) --- CMakeLists.txt | 7 +++ npu/onnx_graph.h | 43 ++++++++++++++ npu/pack_import.cpp | 71 ++++++++++++++++++++++ tests/onnx_graph_test.cpp | 120 ++++++++++++++++++++++++++++++++++++++ 4 files changed, 241 insertions(+) create mode 100644 npu/onnx_graph.h create mode 100644 tests/onnx_graph_test.cpp diff --git a/CMakeLists.txt b/CMakeLists.txt index 1083d3ca9..2cd47cf25 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -1080,6 +1080,10 @@ add_executable(npu_pack_test tests/npu_pack_test.cpp) # (the same reason #1830/#1863 put their arms in a binary CI already compiles). Values come from the encoder reference. add_executable(audio_ops_test tests/audio_ops_test.cpp) target_include_directories(audio_ops_test PRIVATE npu/forward/include) +# The audio pack's graph, read through the engine: pins the topological-order assumption the executor needs and the op +# census the host reference was written against. Skips (77) without a path. +add_executable(onnx_graph_test tests/onnx_graph_test.cpp) +target_link_libraries(onnx_graph_test PRIVATE onebit_npu) 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) @@ -1304,6 +1308,9 @@ if(ONEBIT_NPU_MODEL AND ONEBIT_NPU_BO_DUMP) add_test(NAME npu_pack_model COMMAND npu_pack_test ${ONEBIT_NPU_MODEL} ${ONEBIT_NPU_BO_DUMP}) endif() add_test(NAME audio_ops COMMAND audio_ops_test) +set_tests_properties(audio_ops PROPERTIES SKIP_RETURN_CODE 77) +add_test(NAME onnx_graph COMMAND onnx_graph_test ${ONEBIT_AUDIO_ONNX}) +set_tests_properties(onnx_graph PROPERTIES SKIP_RETURN_CODE 77) # The private NPU route registry (npu/private_route.h), without an add-on. add_executable(npu_private_route_test tests/npu_private_route_test.cpp) target_link_libraries(npu_private_route_test PRIVATE onebit_npu) diff --git a/npu/onnx_graph.h b/npu/onnx_graph.h new file mode 100644 index 000000000..1e6e715fc --- /dev/null +++ b/npu/onnx_graph.h @@ -0,0 +1,43 @@ +// Copyright 2026 bong-water-water-bong +// SPDX-License-Identifier: Apache-2.0 +// +// onnx_graph.h -- a *reader* of an ONNX graph, as opposed to the census walk in npu/pack_import.cpp, which keeps only +// the MatMul-family nodes because that is all a quantised LLM pack needs. An audio pack (the Whisper encoder) is a plain +// f32 graph: 32 transformer blocks of Conv/MatMul/Add/LayerNormalization/Gelu/Softmax/Transpose/Mul, and to execute it the +// engine needs every node in order, with its inputs, outputs and integer attributes, and the set of initializer names. +// +// Nothing here executes anything. It exists so the executor's ordering assumption -- "a node's inputs are graph inputs, +// initializers, or outputs of an earlier node" -- is a *checked* fact about a real pack rather than a hope, and so a pack +// that breaks it is refused by name. + +#pragma once + +#include +#include +#include +#include + +namespace onebit::npu::pack_import { + +struct OnnxNodeInfo { + std::string op; + std::string name; + std::vector ins; + std::vector outs; + std::map attrs; // integer attributes only (axis, perm as a list is not an int attr) + std::map> attrs_list; // `perm`, `pads`, `strides`, `dilations` + std::map attrs_float; // `alpha`, `epsilon` (as stored) +}; + +struct OnnxGraphInfo { + std::vector inputs; // graph inputs, in declaration order + std::vector outputs; // graph outputs + std::vector nodes; // graph order, which ONNX guarantees is topological + std::vector initializers; // weight/bias names + std::vector external; // initializers whose data is in a side file (location set) +}; + +// Throws std::runtime_error with the path and the protobuf complaint on a malformed file. +OnnxGraphInfo parse_onnx_graph(const std::string& path); + +} // namespace onebit::npu::pack_import diff --git a/npu/pack_import.cpp b/npu/pack_import.cpp index 3113f320b..50b68ee87 100644 --- a/npu/pack_import.cpp +++ b/npu/pack_import.cpp @@ -15,6 +15,8 @@ #include "pack_import.h" +#include "onnx_graph.h" + #include "file_map.h" #include "pack.h" #include "q4nx.h" @@ -81,6 +83,9 @@ std::vector pb(const uint8_t* p, size_t n) { } case 5: if (size_t(end - p) < 4) throw std::runtime_error("protobuf: truncated fixed32"); + // engine#1931: keep the four bytes in v. A float attribute (LayerNormalization's epsilon) is wire 5, and + // without this every such attribute reads as 0.0 -- indistinguishable from an absent one. + f.v = uint32_t(p[0]) | (uint32_t(p[1]) << 8) | (uint32_t(p[2]) << 16) | (uint32_t(p[3]) << 24); p += 4; break; default: throw std::runtime_error("protobuf: unsupported wire type " + std::to_string(f.wire)); @@ -4567,4 +4572,70 @@ bool canonical_copy_reusable(const std::string& work, int64_t expected_len) { return !ec && have == static_cast(expected_len); } +// engine#1931: a full graph reader, for packs that are plain f32 graphs (the Whisper encoder) rather than quantised LLM +// packs. The census walk above keeps only the MatMul family by design; this one keeps every node in graph order, which ONNX +// guarantees is topological, plus the initializer names -- so an executor can tell "a weight" from "a graph input", and a +// pack that breaks the ordering is refused by name instead of producing plausible activations. +OnnxGraphInfo parse_onnx_graph(const std::string& path) { + size_t size = 0; + std::string why; + const uint8_t* base = map_file(path, size, &why); + if (base == nullptr) throw std::runtime_error(path + ": " + why); + const auto model = pb(base, size); // ModelProto + const auto* gfield = field1(model, 7); // ModelProto.graph + if (gfield == nullptr) throw std::runtime_error(path + ": no ModelProto.graph (field 7)"); + const auto g = pb(gfield->p, gfield->n); + + OnnxGraphInfo out; + for (const auto* f : fields(g, 11)) { // graph.input + const auto vi = pb(f->p, f->n); + if (const auto* nm = field1(vi, 1)) out.inputs.push_back(pb_str(*nm)); + } + for (const auto* f : fields(g, 12)) { // graph.output + const auto vi = pb(f->p, f->n); + if (const auto* nm = field1(vi, 1)) out.outputs.push_back(pb_str(*nm)); + } + for (const auto* f : fields(g, 5)) { // graph.initializer + const auto t = pb(f->p, f->n); + const auto* nm = field1(t, 8); + if (nm == nullptr) continue; + out.initializers.push_back(pb_str(*nm)); + for (const auto* d : fields(t, 13)) { // TensorProto.external_data: "location" means a side file + const auto kv = pb(d->p, d->n); + const auto* k = field1(kv, 1); + if (k != nullptr && pb_str(*k) == "location") { + out.external.push_back(pb_str(*nm)); + break; + } + } + } + for (const auto* f : fields(g, 1)) { // graph.node + const auto node = pb(f->p, f->n); + OnnxNodeInfo ni; + if (const auto* o = field1(node, 4)) ni.op = pb_str(*o); + if (const auto* nm = field1(node, 3)) ni.name = pb_str(*nm); + for (const auto* i : fields(node, 1)) ni.ins.push_back(pb_str(*i)); + for (const auto* o : fields(node, 2)) ni.outs.push_back(pb_str(*o)); + for (const auto* a : fields(node, 5)) { // AttributeProto + const auto at = pb(a->p, a->n); + const auto* nm = field1(at, 1); + if (nm == nullptr) continue; + const std::string key = pb_str(*nm); + if (const auto* iv = field1(at, 3)) ni.attrs[key] = int64_t(iv->v); // INT + if (const auto* fv = field1(at, 2)) { // FLOAT (wire 5, see fields()) + const uint32_t bits = uint32_t(fv->v); + float f32 = 0.0f; + std::memcpy(&f32, &bits, sizeof f32); + ni.attrs_float[key] = f32; + } + std::vector list; + for (const auto* iv : fields(at, 8)) list.push_back(int64_t(iv->v)); // INTS + if (!list.empty()) ni.attrs_list[key] = std::move(list); + } + out.nodes.push_back(std::move(ni)); + } + unmap_file(base, size); + return out; +} + } // namespace onebit::npu::pack_import diff --git a/tests/onnx_graph_test.cpp b/tests/onnx_graph_test.cpp new file mode 100644 index 000000000..c2a2724a4 --- /dev/null +++ b/tests/onnx_graph_test.cpp @@ -0,0 +1,120 @@ +// Copyright 2026 bong-water-water-bong +// SPDX-License-Identifier: Apache-2.0 +// +// onnx_graph_test.cpp -- the audio pack's graph, read through the engine and checked for the property the executor will +// depend on: every node's inputs are graph inputs, initializers, or outputs of an *earlier* node. Skips (77) without +// ONEBIT_TEST_WHISPER naming the encoder's .onnx. + +#include +#include +#include +#include +#include +#include + +#include "onnx_graph.h" + +namespace { + +int g_failures = 0; + +void check(bool cond, const std::string& what) { + if (!cond) { + std::printf("FAIL %s\n", what.c_str()); + ++g_failures; + } +} + +} // namespace + +int main(int argc, char** argv) { + const char* path = (argc > 1) ? argv[1] : std::getenv("ONEBIT_TEST_WHISPER"); + if (path == nullptr || *path == '\0') { + std::printf("SKIP: pass the encoder .onnx, or set ONEBIT_TEST_WHISPER to it\n"); + return 77; + } + onebit::npu::pack_import::OnnxGraphInfo g; + try { + g = onebit::npu::pack_import::parse_onnx_graph(path); + } catch (const std::exception& e) { + std::printf("FAIL cannot parse %s: %s\n", path, e.what()); + return 1; + } + + check(!g.nodes.empty(), "graph has nodes"); + check(g.inputs.size() == 1, "one graph input (the mel window)"); + check(!g.outputs.empty(), "graph declares an output"); + check(!g.initializers.empty(), "graph has initializers"); + + // The ordering assumption, checked against the real pack rather than assumed. + std::set known(g.inputs.begin(), g.inputs.end()); + known.insert(g.initializers.begin(), g.initializers.end()); + int unnamed_inputs = 0; + for (const auto& n : g.nodes) { + for (const auto& in : n.ins) { + if (in.empty()) continue; // optional input, ONNX spells it "" + if (!known.count(in)) { + if (unnamed_inputs < 5) std::printf("FAIL %s input '%s' is not produced earlier\n", n.op.c_str(), in.c_str()); + ++unnamed_inputs; + } + } + for (const auto& o : n.outs) known.insert(o); + } + check(unnamed_inputs == 0, "every node input is a graph input, an initializer, or an earlier node's output"); + + // The op census the reference implementation was written against (tools/whisper_encoder_ref.py): the front-end is + // exactly two Convs, and a transformer block is MatMul/LayerNormalization/Gelu dominated. + std::map census; + for (const auto& n : g.nodes) ++census[n.op]; + std::printf("nodes %zu, initializers %zu (external %zu), graph input '%s', output '%s'\n", g.nodes.size(), + g.initializers.size(), g.external.size(), g.inputs.empty() ? "" : g.inputs[0].c_str(), + g.outputs.empty() ? "" : g.outputs[0].c_str()); + for (const auto& kv : census) std::printf(" %-24s %d\n", kv.first.c_str(), kv.second); + + // The structure, read off the graph rather than assumed from the reference. Measured on whisper-small's encoder: 447 + // nodes, and three facts that changed how the executor has to be written -- + // * the GELU is NOT a `Gelu` node: it arrives decomposed as Div -> Erf -> Add -> Mul -> Mul, so an executor keyed on + // op names alone needs `Erf` and the surrounding arithmetic (one Erf per block plus two in the front-end); + // * attention is explicit (Transpose/MatMul/Softmax), with 8 MatMuls per layer: q,k,v,o, fc1, fc2, qk, and the + // scores-times-values product; + // * the weights arrive transposed (61 Transpose nodes), which is why the projection shapes in the graph are (K, N). + // The counts scale with depth, so they are checked as ratios against the Softmax count (= one per layer), not as + // absolute constants -- the same test then holds for whisper-large-v3's 32-layer encoder. + check(census["Conv"] == 2, "two Convs (the mel front-end: stride 1 then stride 2, both pad 1)"); + const int layers = census["Softmax"]; + check(layers > 0, "attention scores go through a Softmax"); + if (layers > 0) { + check(census["MatMul"] == 8 * layers, "8 MatMuls per layer (q,k,v,o, fc1, fc2, qk, scores-times-values)"); + check(census["LayerNormalization"] == 2 * layers + 1, "two layer norms per layer plus the encoder's final one"); + // Two exports of this family spell the same activation two ways: whisper-small's encoder_Model decomposes it into + // Div -> Erf -> Add -> Mul -> Mul (0 Gelu nodes, 14 Erf for 12 layers), while whisper-large-v3's carries a real + // `Gelu` op (34 for 32 layers). An executor keyed on op names must accept BOTH; the count, one per layer plus two + // in the front-end, is the invariant that holds either way. + check(census["Erf"] + census["Gelu"] == layers + 2, + "the activation, however spelled, appears once per layer plus twice in the front-end"); + check(census["Erf"] == 0 || census["Gelu"] == 0, "an export uses one spelling of the activation, not a mix"); + } + check(census["Transpose"] > layers, "the weights arrive transposed, so the graph's MatMuls are (K, N)"); + + // One attribute the executor must not guess: LayerNormalization's epsilon is a float attribute, and a reader that + // drops fixed32 (as the census walk does) sees every one of them as absent. + const onebit::npu::pack_import::OnnxNodeInfo* ln = nullptr; + for (const auto& n : g.nodes) + if (n.op == "LayerNormalization") { + ln = &n; + break; + } + if (ln != nullptr) { + const auto it = ln->attrs_float.find("epsilon"); + check(it != ln->attrs_float.end(), "LayerNormalization carries epsilon as a float attribute"); + if (it != ln->attrs_float.end()) check(it->second > 0.0f && it->second < 1e-3f, "epsilon is in the usual range"); + } + + if (g_failures != 0) { + std::printf("onnx graph: %d failures\n", g_failures); + return 1; + } + std::printf("onnx graph: topological order, op census (%d layers, activation spelled %s) and layer-norm epsilon all hold\n", + census["Softmax"], census["Gelu"] ? "Gelu" : "Div/Erf/Add/Mul/Mul"); + return 0; +} From 4bc1d6ed7c738977cd274e66d55d803c5a88f444 Mon Sep 17 00:00:00 2001 From: bong-water-water-bong Date: Tue, 10 Nov 2026 10:43:59 -0400 Subject: [PATCH 14/20] npu: run an audio pack encoder on the host, verified against onnxruntime (engine#1931) --- CMakeLists.txt | 6 +- npu/audio_exec.cpp | 364 ++++++++++++++++++++++++++++++++++++++ npu/audio_exec.h | 32 ++++ npu/onnx_graph.h | 14 ++ npu/pack_import.cpp | 28 ++- tests/audio_exec_test.cpp | 117 ++++++++++++ 6 files changed, 554 insertions(+), 7 deletions(-) create mode 100644 npu/audio_exec.cpp create mode 100644 npu/audio_exec.h create mode 100644 tests/audio_exec_test.cpp diff --git a/CMakeLists.txt b/CMakeLists.txt index 2cd47cf25..dfd0c1c7d 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -566,7 +566,7 @@ if(ONEBIT_NPU_HOST) find_package(PkgConfig REQUIRED) pkg_check_modules(PCRE2 REQUIRED IMPORTED_TARGET libpcre2-8) add_library(onebit_npu STATIC npu/file_map.cpp npu/flm_head_import.cpp npu/full_elf.cpp npu/model.cpp - npu/moe.cpp npu/pack.cpp npu/pack_import.cpp npu/private_route.cpp + npu/audio_exec.cpp npu/moe.cpp npu/pack.cpp npu/pack_import.cpp npu/private_route.cpp npu/q4nx.cpp npu/q4nx_device.cpp npu/tokenizer.cpp) target_include_directories(onebit_npu PUBLIC npu) target_link_libraries(onebit_npu PUBLIC nlohmann_json::nlohmann_json PRIVATE PkgConfig::PCRE2) @@ -1084,6 +1084,8 @@ target_include_directories(audio_ops_test PRIVATE npu/forward/include) # census the host reference was written against. Skips (77) without a path. add_executable(onnx_graph_test tests/onnx_graph_test.cpp) target_link_libraries(onnx_graph_test PRIVATE onebit_npu) +add_executable(audio_exec_test tests/audio_exec_test.cpp) +target_link_libraries(audio_exec_test PRIVATE onebit_npu) 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) @@ -1310,6 +1312,8 @@ endif() add_test(NAME audio_ops COMMAND audio_ops_test) set_tests_properties(audio_ops PROPERTIES SKIP_RETURN_CODE 77) add_test(NAME onnx_graph COMMAND onnx_graph_test ${ONEBIT_AUDIO_ONNX}) +add_test(NAME audio_exec COMMAND audio_exec_test ${ONEBIT_AUDIO_ONNX}) +set_tests_properties(audio_exec PROPERTIES SKIP_RETURN_CODE 77) set_tests_properties(onnx_graph PROPERTIES SKIP_RETURN_CODE 77) # The private NPU route registry (npu/private_route.h), without an add-on. add_executable(npu_private_route_test tests/npu_private_route_test.cpp) diff --git a/npu/audio_exec.cpp b/npu/audio_exec.cpp new file mode 100644 index 000000000..3f7f28015 --- /dev/null +++ b/npu/audio_exec.cpp @@ -0,0 +1,364 @@ +// Copyright 2026 bong-water-water-bong +// SPDX-License-Identifier: Apache-2.0 +// +// audio_exec.cpp -- the host f32 executor for an audio pack's graph (see audio_exec.h for the contract). + +#include "audio_exec.h" + +#include +#include +#include +#include +#include + +#include "forward/include/audio_ops.h" +#include "file_map.h" + +namespace onebit::npu::pack_import { +namespace { + +struct Tensor { + std::vector dims; + std::vector data; +}; + +size_t nelem(const std::vector& d) { + size_t n = 1; + for (int64_t v : d) n *= size_t(v < 0 ? 0 : v); + return n; +} + +[[noreturn]] void fail(const std::string& why) { throw std::runtime_error(why); } + +// numpy-style right-aligned broadcast of `small` against `big`; returns the result dims and, per output index, the source +// index. Only the two forms this graph uses: equal, or `small` has trailing dims of 1/size and fewer ranks. +struct Broadcast { + std::vector dims; + std::vector stride_big, stride_small; // 0 where the axis is broadcast +}; + +Broadcast plan_broadcast(const std::vector& a, const std::vector& b) { + Broadcast p; + const size_t r = std::max(a.size(), b.size()); + p.dims.assign(r, 1); + p.stride_big.assign(r, 0); + p.stride_small.assign(r, 0); + for (size_t i = 0; i < r; ++i) { + const int64_t da = (i < r - a.size()) ? 1 : a[i - (r - a.size())]; + const int64_t db = (i < r - b.size()) ? 1 : b[i - (r - b.size())]; + if (da != db && da != 1 && db != 1) fail("broadcast mismatch " + std::to_string(da) + " vs " + std::to_string(db)); + p.dims[i] = da == 1 ? db : da; + } + size_t sb = 1, ss = 1; + for (size_t i = r; i-- > 0;) { + const int64_t da = (i < r - a.size()) ? 1 : a[i - (r - a.size())]; + const int64_t db = (i < r - b.size()) ? 1 : b[i - (r - b.size())]; + p.stride_big[i] = da == 1 ? 0 : sb; + p.stride_small[i] = db == 1 ? 0 : ss; + sb *= size_t(da); + ss *= size_t(db); + } + return p; +} + +void broadcast_binary(const Tensor& a, const Tensor& b, const std::string& op, Tensor* out) { + const Broadcast p = plan_broadcast(a.dims, b.dims); + out->dims = p.dims; + out->data.resize(nelem(p.dims)); + std::vector idx(p.dims.size(), 0); + for (size_t o = 0; o < out->data.size(); ++o) { + size_t ia = 0, ib = 0; + for (size_t d = 0; d < idx.size(); ++d) { + ia += idx[d] * p.stride_big[d]; + ib += idx[d] * p.stride_small[d]; + } + const float x = a.data[ia], y = b.data[ib]; + out->data[o] = (op == "Add") ? x + y : (op == "Sub") ? x - y : (op == "Mul") ? x * y : x / y; + for (size_t d = idx.size(); d-- > 0;) { + if (++idx[d] < size_t(p.dims[d])) break; + idx[d] = 0; + } + } +} + +std::vector dims_of(const OnnxInit& t) { return t.dims; } + +// Copies an initializer to f32, following external_data into the pack's side file when the bytes are not inline. +void load_init_f32(const OnnxGraphInfo& g, const OnnxInit& t, std::vector* dst) { + const uint8_t* src = t.raw; + size_t len = t.raw_len; + size_t size = 0; + std::string why; + const uint8_t* mapped = nullptr; + if (!t.location.empty()) { + mapped = map_file(g.dir + "/" + t.location, size, &why); + if (mapped == nullptr) fail("external_data " + t.location + ": " + why); + const uint64_t off = t.offset, want = t.length; + if (off > size || want > size - off) fail("external_data " + t.location + ": range past the end"); + src = mapped + off; + len = size_t(want); + if (len == 0) len = size_t(size - off); + } + if (src == nullptr) fail("initializer '" + t.name + "' has no data"); + const size_t n = nelem(t.dims); + dst->assign(n, 0.0f); + if (t.dtype == 1) { // FLOAT + if (len < n * 4) fail("initializer '" + t.name + "' is short: " + std::to_string(len) + " < " + std::to_string(n * 4)); + std::memcpy(dst->data(), src, n * 4); + } else if (t.dtype == 7) { // INT64 -- shape operands + if (len < n * 8) fail("int64 initializer '" + t.name + "' is short"); + std::memcpy(dst->data(), src, n * 8); // bit pattern; callers that need the integer read it via load_init_i64 + } else if (t.dtype == 10) { // FLOAT16 + if (len < n * 2) fail("f16 initializer '" + t.name + "' is short"); + for (size_t i = 0; i < n; ++i) { + uint16_t h = 0; + std::memcpy(&h, src + i * 2, 2); + const uint32_t sign = uint32_t(h & 0x8000) << 16, exp = (h >> 10) & 0x1f, man = h & 0x3ff; + uint32_t bits; + if (exp == 0) { + bits = sign | (man ? (uint32_t(127 - 15) << 23) + (man << 13) : 0); + } else if (exp == 31) { + bits = sign | 0x7f800000u | (man << 13); + } else { + bits = sign | ((exp + 127 - 15) << 23) | (man << 13); + } + std::memcpy(&(*dst)[i], &bits, 4); + } + } else { + fail("initializer '" + t.name + "' has dtype " + std::to_string(t.dtype) + ", which this executor does not load"); + } + if (mapped != nullptr) unmap_file(mapped, size); +} + +void load_init_i64(const OnnxGraphInfo& g, const OnnxInit& t, std::vector* dst) { + const uint8_t* src = t.raw; + size_t len = t.raw_len; + size_t size = 0; + std::string why; + const uint8_t* mapped = nullptr; + if (!t.location.empty()) { + mapped = map_file(g.dir + "/" + t.location, size, &why); + if (mapped == nullptr) fail("external_data " + t.location + ": " + why); + if (t.offset > size || t.length > size - t.offset) fail("external_data range past the end"); + src = mapped + t.offset; + len = size_t(t.length); + } + if (t.dtype != 7) fail("'" + t.name + "' is not INT64"); + const size_t n = nelem(t.dims); + if (src == nullptr || len < n * 8) fail("int64 initializer '" + t.name + "' is short"); + dst->assign(n, 0); + std::memcpy(dst->data(), src, n * 8); + if (mapped != nullptr) unmap_file(mapped, size); +} + +} // namespace + +bool run_audio_graph(const OnnxGraphInfo& g, const float* mel, int64_t mel_bins, int64_t frames, const std::string& out_name, + AudioRunResult* out, std::string* why) { + try { + std::map t; + // The graph input: the mel window, in the layout the first Conv consumes. + Tensor in; + in.dims = {1, mel_bins, frames}; + in.data.assign(mel, mel + size_t(mel_bins) * size_t(frames)); + if (g.inputs.empty()) fail("graph declares no input"); + t[g.inputs[0]] = std::move(in); + + std::map inits; + for (const auto& ti : g.init_tensors) inits[ti.name] = &ti; + + auto need = [&](const std::string& name) -> const Tensor& { + const auto it = t.find(name); + if (it == t.end()) fail("node input '" + name + "' was never produced"); + return it->second; + }; + auto init_f32 = [&](const std::string& name) -> const Tensor& { + auto it = t.find(name); + if (it != t.end()) return it->second; + const auto ii = inits.find(name); + if (ii == inits.end()) fail("'" + name + "' is neither a tensor nor an initializer"); + Tensor v; + v.dims = dims_of(*ii->second); + load_init_f32(g, *ii->second, &v.data); + return t.emplace(name, std::move(v)).first->second; + }; + + int node_index = 0; + for (const auto& n : g.nodes) { + ++node_index; + const auto where = [&]() { return "node " + std::to_string(node_index) + " (" + n.op + ")"; }; + if (n.outs.empty()) fail(where() + " has no output"); + const std::string& oname = n.outs[0]; + Tensor r; + if (n.op == "MatMul") { + const Tensor& a = need(n.ins[0]); + const Tensor& b = t.count(n.ins[1]) != 0 ? need(n.ins[1]) : init_f32(n.ins[1]); + if (a.dims.size() < 2 || b.dims.size() < 2) fail(where() + ": needs rank >= 2 on both sides"); + const size_t ra = a.dims.size(), rb = b.dims.size(); + const int64_t M = a.dims[ra - 2], K = a.dims.back(), K2 = b.dims[rb - 2], N = b.dims.back(); + if (K != K2) fail(where() + ": " + std::to_string(K) + " vs " + std::to_string(K2)); + // Attention's q.k^T and scores.v are 4-D: (batch, heads, M|K, K|N). The batch dims broadcast, the last two + // are a plain GEMM, so one projection covers both the layer weights and the attention products. + const std::vector ab(a.dims.begin(), a.dims.end() - 2), bb(b.dims.begin(), b.dims.end() - 2); + const Broadcast bp = plan_broadcast(ab, bb); + r.dims = bp.dims; + r.dims.push_back(M); + r.dims.push_back(N); + r.data.assign(nelem(r.dims), 0.0f); + std::vector idx(bp.dims.size(), 0); + const size_t batches = nelem(bp.dims); + for (size_t bi = 0; bi < batches; ++bi) { + size_t oa = 0, ob = 0; + for (size_t d = 0; d < idx.size(); ++d) { + oa += idx[d] * bp.stride_big[d]; + ob += idx[d] * bp.stride_small[d]; + } + const float* ap = a.data.data() + oa * size_t(M * K); + const float* bpp = b.data.data() + ob * size_t(K * N); + float* rp = r.data.data() + bi * size_t(M * N); + for (int64_t i = 0; i < M; ++i) { + float* rr = rp + size_t(i) * size_t(N); + for (int64_t k = 0; k < K; ++k) { + const float av = ap[size_t(i) * size_t(K) + size_t(k)]; + if (av == 0.0f) continue; + const float* br = bpp + size_t(k) * size_t(N); + for (int64_t j = 0; j < N; ++j) rr[j] += av * br[j]; + } + } + for (size_t d = idx.size(); d-- > 0;) { + if (++idx[d] < size_t(bp.dims[d])) break; + idx[d] = 0; + } + } + } else if (n.op == "Add" || n.op == "Mul" || n.op == "Div" || n.op == "Sub") { + // An operand is either a tensor an earlier node produced or an initializer (a bias vector, the 1/sqrt(d) + // scalar); init_f32() handles the second without a separate code path at every call site. + auto operand = [&](const std::string& name) -> const Tensor& { + return t.count(name) != 0 ? need(name) : init_f32(name); + }; + if (n.ins.size() < 2) fail(where() + ": needs two operands"); + broadcast_binary(operand(n.ins[0]), operand(n.ins[1]), n.op, &r); + } else if (n.op == "Erf") { + const Tensor& a = need(n.ins[0]); + r.dims = a.dims; + r.data.resize(a.data.size()); + for (size_t i = 0; i < a.data.size(); ++i) r.data[i] = std::erf(a.data[i]); + } else if (n.op == "Gelu") { + const Tensor& a = need(n.ins[0]); + r.dims = a.dims; + r.data.resize(a.data.size()); + const auto ap = n.attrs.find("approximate"); + (void)ap; // ONNX spells it as a string; the export here is 'none', which is what audio_ops::gelu_erf is + for (size_t i = 0; i < a.data.size(); ++i) r.data[i] = audio::gelu_erf(a.data[i]); + } else if (n.op == "LayerNormalization") { + const Tensor& a = need(n.ins[0]); + const Tensor& scale = init_f32(n.ins[1]); + const Tensor* bias = (n.ins.size() > 2 && !n.ins[2].empty()) ? &init_f32(n.ins[2]) : nullptr; + const int64_t axis = n.attrs.count("axis") ? n.attrs.at("axis") : -1; + if (axis != -1 && axis != int64_t(a.dims.size()) - 1) fail(where() + ": only the last axis is implemented"); + const int64_t nrm = a.dims.back(); + const float eps = n.attrs_float.count("epsilon") ? n.attrs_float.at("epsilon") : 1e-5f; + if (int64_t(scale.data.size()) != nrm) fail(where() + ": scale is not the normalised width"); + r.dims = a.dims; + r.data.resize(a.data.size()); + for (size_t row = 0; row < a.data.size() / size_t(nrm); ++row) { + audio::layer_norm(a.data.data() + row * size_t(nrm), scale.data.data(), + bias != nullptr ? bias->data.data() : nullptr, int(nrm), eps, + r.data.data() + row * size_t(nrm)); + } + } else if (n.op == "Softmax") { + const Tensor& a = need(n.ins[0]); + const int64_t axis = n.attrs.count("axis") ? n.attrs.at("axis") : -1; + if (axis != -1 && axis != int64_t(a.dims.size()) - 1) fail(where() + ": only the last axis is implemented"); + r = a; + const int64_t nrm = a.dims.back(); + for (size_t row = 0; row < r.data.size() / size_t(nrm); ++row) + audio::softmax_row(r.data.data() + row * size_t(nrm), int(nrm)); + } else if (n.op == "Transpose") { + const Tensor& a = need(n.ins[0]); + const auto it = n.attrs_list.find("perm"); + std::vector perm; + if (it != n.attrs_list.end()) perm = it->second; + else + for (size_t i = a.dims.size(); i-- > 0;) perm.push_back(int64_t(i)); + if (perm.size() != a.dims.size()) fail(where() + ": perm rank mismatch"); + r.dims.resize(perm.size()); + for (size_t i = 0; i < perm.size(); ++i) r.dims[i] = a.dims[size_t(perm[i])]; + r.data.resize(a.data.size()); + std::vector in_stride(perm.size(), 1), out_stride(perm.size(), 1); + for (size_t i = perm.size(); i-- > 1;) { + in_stride[i - 1] = in_stride[i] * size_t(a.dims[i]); + out_stride[i - 1] = out_stride[i] * size_t(r.dims[i]); + } + std::vector idx(perm.size(), 0); + for (size_t o = 0; o < r.data.size(); ++o) { + size_t src = 0; + for (size_t i = 0; i < idx.size(); ++i) src += idx[i] * in_stride[size_t(perm[i])]; + r.data[o] = a.data[src]; + for (size_t d = idx.size(); d-- > 0;) { + if (++idx[d] < size_t(r.dims[d])) break; + idx[d] = 0; + } + } + } else if (n.op == "Reshape") { + const Tensor& a = need(n.ins[0]); + std::vector want; + auto ii = inits.find(n.ins[1]); + if (ii == inits.end()) fail(where() + ": target shape is not an initializer"); + load_init_i64(g, *ii->second, &want); + std::vector dims(want.size()); + int64_t infer = -1; + size_t known = 1; + for (size_t i = 0; i < want.size(); ++i) { + int64_t d = want[i]; + if (d == 0) d = a.dims[i]; + if (d == -1) { + if (infer >= 0) fail(where() + ": two -1"); + infer = int64_t(i); + d = 1; + } + dims[i] = d; + known *= size_t(d); + } + if (infer >= 0) dims[size_t(infer)] = int64_t(a.data.size() / std::max(known, 1)); + if (nelem(dims) != a.data.size()) fail(where() + ": reshape changes the element count"); + r.dims = dims; + r.data = a.data; + } else if (n.op == "Conv") { + const Tensor& a = need(n.ins[0]); + const Tensor& w = init_f32(n.ins[1]); + const Tensor* bias = (n.ins.size() > 2 && !n.ins[2].empty()) ? &init_f32(n.ins[2]) : nullptr; + const auto gs = n.attrs_list.find("strides"); + const auto gp = n.attrs_list.find("pads"); + const int stride = (gs != n.attrs_list.end() && !gs->second.empty()) ? int(gs->second[0]) : 1; + const int pad = (gp != n.attrs_list.end() && !gp->second.empty()) ? int(gp->second[0]) : 0; + if (a.dims.size() != 3 || w.dims.size() != 3) fail(where() + ": only 1-D temporal Conv on (N,C,T)"); + const int64_t cin = a.dims[1], T = a.dims[2]; + const int64_t cout = w.dims[0], k = w.dims[2]; + if (w.dims[1] != cin) fail(where() + ": kernel channels do not match the input"); + const int64_t tout = audio::conv1d_out_len(int(T), int(k), stride, pad); + r.dims = {a.dims[0], cout, tout}; + r.data.resize(size_t(a.dims[0]) * size_t(cout) * size_t(tout)); + for (int64_t n0 = 0; n0 < a.dims[0]; ++n0) + audio::conv1d(a.data.data() + size_t(n0) * size_t(cin) * size_t(T), int(cin), int(T), w.data.data(), + bias != nullptr ? bias->data.data() : nullptr, int(cout), int(k), stride, pad, + r.data.data() + size_t(n0) * size_t(cout) * size_t(tout)); + } else { + fail(where() + " is not implemented; the pack needs an op this executor does not have"); + } + t[oname] = std::move(r); + } + + const std::string want = out_name.empty() ? (g.outputs.empty() ? std::string() : g.outputs[0]) : out_name; + const auto it = t.find(want); + if (it == t.end()) fail("output '" + want + "' was never produced"); + out->dims = it->second.dims; + out->data = it->second.data; + return true; + } catch (const std::exception& e) { + if (why != nullptr) *why = e.what(); + return false; + } +} + +} // namespace onebit::npu::pack_import diff --git a/npu/audio_exec.h b/npu/audio_exec.h new file mode 100644 index 000000000..31ec1ee58 --- /dev/null +++ b/npu/audio_exec.h @@ -0,0 +1,32 @@ +// Copyright 2026 bong-water-water-bong +// SPDX-License-Identifier: Apache-2.0 +// +// audio_exec.h -- run an audio pack's graph on the host in f32. Not a general ONNX runtime: it covers exactly the ops a +// Whisper encoder uses (measured on both exports -- 447 nodes / 11 op types for whisper-small, 1031 / 10 for large-v3), and +// it refuses an op it does not implement *by name* rather than producing a plausible activation from a silently skipped node. +// +// The four f32 kernels this stands on are in npu/forward/include/audio_ops.h, which is device-free and known-answer tested; +// this file adds the shape bookkeeping and the wiring. The caller compares its output against the reference, and against +// onnxruntime, before anything here is used to argue a device number. + +#pragma once + +#include +#include + +#include "onnx_graph.h" + +namespace onebit::npu::pack_import { + +struct AudioRunResult { + std::vector dims; // dims of the requested output + std::vector data; // row-major +}; + +// mel: (mel_bins, frames) row-major f32, the encoder's graph input in the layout its own first Conv consumes. +// out_name: the graph output to return (empty = the graph's declared output). +// Returns true and fills `out`; on failure returns false and sets `why` to the node and the reason. +bool run_audio_graph(const OnnxGraphInfo& g, const float* mel, int64_t mel_bins, int64_t frames, + const std::string& out_name, AudioRunResult* out, std::string* why); + +} // namespace onebit::npu::pack_import diff --git a/npu/onnx_graph.h b/npu/onnx_graph.h index 1e6e715fc..db3a8157a 100644 --- a/npu/onnx_graph.h +++ b/npu/onnx_graph.h @@ -29,7 +29,21 @@ struct OnnxNodeInfo { std::map attrs_float; // `alpha`, `epsilon` (as stored) }; +// One initializer: its shape, dtype and where its bytes are (inline `raw_data`, or a side file via external_data). +struct OnnxInit { + std::string name; + int dtype = 0; // TensorProto.DataType: 1 = FLOAT, 7 = INT64, 10 = FLOAT16, ... + std::vector dims; + const uint8_t* raw = nullptr; // valid only while the mapped file is open; the caller copies what it needs + size_t raw_len = 0; + std::string location; // external_data side file, relative to the graph's directory + uint64_t offset = 0; + uint64_t length = 0; +}; + struct OnnxGraphInfo { + std::string dir; // directory of the .onnx file, for resolving external_data + std::vector init_tensors; std::vector inputs; // graph inputs, in declaration order std::vector outputs; // graph outputs std::vector nodes; // graph order, which ONNX guarantees is topological diff --git a/npu/pack_import.cpp b/npu/pack_import.cpp index 50b68ee87..c9c298d06 100644 --- a/npu/pack_import.cpp +++ b/npu/pack_import.cpp @@ -4587,6 +4587,8 @@ OnnxGraphInfo parse_onnx_graph(const std::string& path) { const auto g = pb(gfield->p, gfield->n); OnnxGraphInfo out; + const size_t slash = path.find_last_of('/'); + out.dir = (slash == std::string::npos) ? std::string(".") : path.substr(0, slash); for (const auto* f : fields(g, 11)) { // graph.input const auto vi = pb(f->p, f->n); if (const auto* nm = field1(vi, 1)) out.inputs.push_back(pb_str(*nm)); @@ -4600,14 +4602,26 @@ OnnxGraphInfo parse_onnx_graph(const std::string& path) { const auto* nm = field1(t, 8); if (nm == nullptr) continue; out.initializers.push_back(pb_str(*nm)); - for (const auto* d : fields(t, 13)) { // TensorProto.external_data: "location" means a side file + OnnxInit ti; + ti.name = pb_str(*nm); + if (const auto* dt = field1(t, 2)) ti.dtype = int(dt->v); + for (const auto* d : fields(t, 1)) ti.dims.push_back(int64_t(d->v)); + if (const auto* rd = field1(t, 9)) { + ti.raw = rd->p; + ti.raw_len = rd->n; + } + for (const auto* d : fields(t, 13)) { // TensorProto.external_data: key/value string pairs const auto kv = pb(d->p, d->n); const auto* k = field1(kv, 1); - if (k != nullptr && pb_str(*k) == "location") { - out.external.push_back(pb_str(*nm)); - break; - } + const auto* v = field1(kv, 2); + if (k == nullptr || v == nullptr) continue; + const std::string key = pb_str(*k), val = pb_str(*v); + if (key == "location") ti.location = val; + else if (key == "offset") ti.offset = std::stoull(val); + else if (key == "length") ti.length = std::stoull(val); } + if (!ti.location.empty()) out.external.push_back(ti.name); + out.init_tensors.push_back(std::move(ti)); } for (const auto* f : fields(g, 1)) { // graph.node const auto node = pb(f->p, f->n); @@ -4634,7 +4648,9 @@ OnnxGraphInfo parse_onnx_graph(const std::string& path) { } out.nodes.push_back(std::move(ni)); } - unmap_file(base, size); + // NOT unmapped: init_tensors[].raw points into this mapping, and the executor copies what it needs while it is alive. + // shortcut: the mapping is released when the process exits; a long-lived caller that parses many graphs should own the + // mapping itself, which is the refactor to make when a second consumer appears. return out; } diff --git a/tests/audio_exec_test.cpp b/tests/audio_exec_test.cpp new file mode 100644 index 000000000..4d5c163ef --- /dev/null +++ b/tests/audio_exec_test.cpp @@ -0,0 +1,117 @@ +// Copyright 2026 bong-water-water-bong +// SPDX-License-Identifier: Apache-2.0 +// +// audio_exec_test.cpp -- run an audio pack's encoder through the engine's own host executor and compare it against +// onnxruntime, on an input both sides can build from a formula (no binary fixture in the repo). +// +// The mel is mel[b,t] = sin(0.013*(b+1)*(t+1))*0.5 + 0.25 for 80 bins and 3000 frames, computed in double on both sides. +// Expected values below are onnxruntime's, on whisper-small's encoder_model.onnx. Skips (77) without the pack. + +#include +#include +#include +#include +#include + +#include "audio_exec.h" +#include "onnx_graph.h" + +namespace { + +int g_failures = 0; + +void check(bool cond, const std::string& what) { + if (!cond) { + std::printf("FAIL %s\n", what.c_str()); + ++g_failures; + } +} + +void expect_close(float got, double want, double tol, const std::string& what) { + const double d = std::fabs(double(got) - want); + if (!(d <= tol)) { + std::printf("FAIL %s: got %.7g want %.7g (|d| %.3g)\n", what.c_str(), double(got), want, d); + ++g_failures; + } +} + +} // namespace + +int main(int argc, char** argv) { + const char* path = (argc > 1) ? argv[1] : std::getenv("ONEBIT_TEST_WHISPER"); + if (path == nullptr || *path == '\0') { + std::printf("SKIP: pass the encoder .onnx, or set ONEBIT_TEST_WHISPER to it\n"); + return 77; + } + using namespace onebit::npu::pack_import; + OnnxGraphInfo g; + try { + g = parse_onnx_graph(path); + } catch (const std::exception& e) { + std::printf("FAIL cannot parse %s: %s\n", path, e.what()); + return 1; + } + + // The mel-bin count is the front-end Conv's input-channel count: 80 for whisper-small, 128 for large-v3. Read from the + // pack rather than hard-coded, so the same test runs on either (the numeric comparison below is small's). + int bins = 80; + for (const auto& ti : g.init_tensors) + if (ti.dims.size() == 3 && ti.dims[0] == 768) { bins = int(ti.dims[1]); break; } + for (const auto& n : g.nodes) + if (n.op == "Conv") { + for (const auto& ti : g.init_tensors) + if (ti.name == n.ins[1] && ti.dims.size() == 3) bins = int(ti.dims[1]); + break; + } + const int frames = 3000; + std::vector mel(size_t(bins) * frames); + for (int b = 0; b < bins; ++b) + for (int t = 0; t < frames; ++t) + mel[size_t(b) * frames + t] = float(std::sin(0.013 * (b + 1) * (t + 1)) * 0.5 + 0.25); + + AudioRunResult r; + std::string why; + if (!run_audio_graph(g, mel.data(), bins, frames, "", &r, &why)) { + std::printf("FAIL run: %s\n", why.c_str()); + return 1; + } + std::printf("ran %zu nodes -> dims", g.nodes.size()); + for (int64_t d : r.dims) std::printf(" %lld", (long long)d); + std::printf("\n"); + check(r.dims.size() == 3 && r.dims[0] == 1 && r.dims[1] == 1500, "encoder output is [1,1500,hidden]"); + if (r.dims.size() != 3 || r.dims[2] != 768) { + // A larger encoder: the point of this run is that the graph executed (including external_data weights), not the + // small model's numbers. + std::printf("audio exec: ran a %d-bin encoder of hidden %lld (numeric comparison is whisper-small's only)\n", bins, + (long long)(r.dims.size() == 3 ? r.dims[2] : 0)); + return g_failures == 0 ? 0 : 1; + } + + const int64_t H = 768; + auto at = [&](int64_t t, int64_t h) { return r.data[size_t(t) * H + size_t(h)]; }; + double mean = 0.0, var = 0.0; + for (float v : r.data) mean += v; + mean /= double(r.data.size()); + for (float v : r.data) var += (double(v) - mean) * (double(v) - mean); + var /= double(r.data.size()); + std::printf("mean %.9f std %.9f\n", mean, std::sqrt(var)); + + // onnxruntime, same input, same pack: mean -0.009045244, std 1.678493023 + expect_close(float(mean), -0.009045244, 2e-4, "output mean"); + expect_close(float(std::sqrt(var)), 1.678493023, 2e-4, "output std"); + + // frame 0, 750 and 1499 -- the first frame, the middle, and the last (whose window is padding-heavy) + const double p0[5] = {-0.5922473669, -1.6559172869, 0.0451035500, -0.9414483309, -1.6385756731}; + const double p1[4] = {0.1923044026, -0.0132405162, 0.4830040932, 0.6385290623}; + const double p2[4] = {-2.1586399078, 0.3378775120, 0.8375968933, 0.5081450939}; + for (int i = 0; i < 5; ++i) expect_close(at(0, i), p0[i], 2e-3, "frame 0 col " + std::to_string(i)); + for (int i = 0; i < 4; ++i) expect_close(at(750, i), p1[i], 2e-3, "frame 750 col " + std::to_string(i)); + for (int i = 0; i < 4; ++i) expect_close(at(1499, i), p2[i], 2e-3, "frame 1499 col " + std::to_string(i)); + + if (g_failures != 0) { + std::printf("audio exec: %d failures\n", g_failures); + return 1; + } + std::printf("audio exec: whole encoder on the host matches onnxruntime\n"); + return 0; +} From 28a6c62f9eb7cbf07bc5a1441afc2f656aeddeda Mon Sep 17 00:00:00 2001 From: bong-water-water-bong Date: Tue, 10 Nov 2026 11:11:55 -0400 Subject: [PATCH 15/20] npu: run the Whisper decoder too -- 1233 nodes, both graphs, verified against onnxruntime (engine#1931) --- npu/audio_exec.cpp | 242 +++++++++++++++++++++++++++++++++++--- npu/audio_exec.h | 15 ++- tests/audio_exec_test.cpp | 45 +++++++ 3 files changed, 282 insertions(+), 20 deletions(-) diff --git a/npu/audio_exec.cpp b/npu/audio_exec.cpp index 3f7f28015..a2b73c7e5 100644 --- a/npu/audio_exec.cpp +++ b/npu/audio_exec.cpp @@ -9,6 +9,7 @@ #include #include #include +#include #include #include "forward/include/audio_ops.h" @@ -19,7 +20,10 @@ namespace { struct Tensor { std::vector dims; - std::vector data; + std::vector data; // f32 payload + std::vector i64; // int64 payload (tokens, indices, shapes); exactly one of the two is used + std::vector b; // bool payload (a mask) + size_t count() const { return i64.empty() ? (b.empty() ? data.size() : b.size()) : i64.size(); } }; size_t nelem(const std::vector& d) { @@ -61,6 +65,21 @@ Broadcast plan_broadcast(const std::vector& a, const std::vector& od, const std::vector& src, const std::vector& idx) { + size_t off = 0, stride = 1; + for (size_t i = src.size(); i-- > 0;) { + const size_t from_right = src.size() - i; // 1 = last axis + const size_t oi = od.size() - from_right; // the output axis it aligns with + const int64_t d = src[i]; + const size_t k = (oi < od.size() && d != 1) ? idx[oi] : 0; + off += k * stride; + stride *= size_t(d); + } + return off; +} + void broadcast_binary(const Tensor& a, const Tensor& b, const std::string& op, Tensor* out) { const Broadcast p = plan_broadcast(a.dims, b.dims); out->dims = p.dims; @@ -84,7 +103,7 @@ void broadcast_binary(const Tensor& a, const Tensor& b, const std::string& op, T std::vector dims_of(const OnnxInit& t) { return t.dims; } // Copies an initializer to f32, following external_data into the pack's side file when the bytes are not inline. -void load_init_f32(const OnnxGraphInfo& g, const OnnxInit& t, std::vector* dst) { +void load_init_f32(const OnnxGraphInfo& g, const OnnxInit& t, Tensor* dst) { const uint8_t* src = t.raw; size_t len = t.raw_len; size_t size = 0; @@ -101,13 +120,17 @@ void load_init_f32(const OnnxGraphInfo& g, const OnnxInit& t, std::vector } if (src == nullptr) fail("initializer '" + t.name + "' has no data"); const size_t n = nelem(t.dims); - dst->assign(n, 0.0f); + if (t.dtype == 7) { // INT64 -- indices and shape operands; the payload stays integer + dst->i64.assign(n, 0); + if (len < n * 8) fail("int64 initializer '" + t.name + "' is short"); + std::memcpy(dst->i64.data(), src, n * 8); + if (mapped != nullptr) unmap_file(mapped, size); + return; + } + dst->data.assign(n, 0.0f); if (t.dtype == 1) { // FLOAT if (len < n * 4) fail("initializer '" + t.name + "' is short: " + std::to_string(len) + " < " + std::to_string(n * 4)); - std::memcpy(dst->data(), src, n * 4); - } else if (t.dtype == 7) { // INT64 -- shape operands - if (len < n * 8) fail("int64 initializer '" + t.name + "' is short"); - std::memcpy(dst->data(), src, n * 8); // bit pattern; callers that need the integer read it via load_init_i64 + std::memcpy(dst->data.data(), src, n * 4); } else if (t.dtype == 10) { // FLOAT16 if (len < n * 2) fail("f16 initializer '" + t.name + "' is short"); for (size_t i = 0; i < n; ++i) { @@ -122,7 +145,7 @@ void load_init_f32(const OnnxGraphInfo& g, const OnnxInit& t, std::vector } else { bits = sign | ((exp + 127 - 15) << 23) | (man << 13); } - std::memcpy(&(*dst)[i], &bits, 4); + std::memcpy(&dst->data[i], &bits, 4); } } else { fail("initializer '" + t.name + "' has dtype " + std::to_string(t.dtype) + ", which this executor does not load"); @@ -153,16 +176,11 @@ void load_init_i64(const OnnxGraphInfo& g, const OnnxInit& t, std::vector t, const std::string& out_name, AudioRunResult* out, + std::string* why) { try { - std::map t; - // The graph input: the mel window, in the layout the first Conv consumes. - Tensor in; - in.dims = {1, mel_bins, frames}; - in.data.assign(mel, mel + size_t(mel_bins) * size_t(frames)); - if (g.inputs.empty()) fail("graph declares no input"); - t[g.inputs[0]] = std::move(in); std::map inits; for (const auto& ti : g.init_tensors) inits[ti.name] = &ti; @@ -179,10 +197,16 @@ bool run_audio_graph(const OnnxGraphInfo& g, const float* mel, int64_t mel_bins, if (ii == inits.end()) fail("'" + name + "' is neither a tensor nor an initializer"); Tensor v; v.dims = dims_of(*ii->second); - load_init_f32(g, *ii->second, &v.data); + load_init_f32(g, *ii->second, &v); return t.emplace(name, std::move(v)).first->second; }; + // A node operand is either a tensor an earlier node produced or an initializer (a weight, an index vector, a + // scalar attribute operand). `any_t` is the one spelling of that rule. + auto any_t = [&](const std::string& name) -> const Tensor& { + return t.count(name) != 0 ? need(name) : init_f32(name); + }; + int node_index = 0; for (const auto& n : g.nodes) { ++node_index; @@ -343,6 +367,145 @@ bool run_audio_graph(const OnnxGraphInfo& g, const float* mel, int64_t mel_bins, audio::conv1d(a.data.data() + size_t(n0) * size_t(cin) * size_t(T), int(cin), int(T), w.data.data(), bias != nullptr ? bias->data.data() : nullptr, int(cout), int(k), stride, pad, r.data.data() + size_t(n0) * size_t(cout) * size_t(tout)); + } else if (n.op == "Gather") { + // The token embedding: axis 0, indices are the int64 token window. + const Tensor& a = any_t(n.ins[0]); + const Tensor& ix = any_t(n.ins[1]); + if (!ix.i64.empty()) { + const int64_t axis = n.attrs.count("axis") ? n.attrs.at("axis") : 0; + if (axis != 0 || a.dims.empty()) fail(where() + ": only axis 0 is implemented"); + const int64_t rows = a.dims[0]; + const size_t row = a.data.size() / size_t(rows); + r.dims = ix.dims; + for (size_t i = 1; i < a.dims.size(); ++i) r.dims.push_back(a.dims[i]); + r.data.resize(ix.i64.size() * row); + for (size_t i = 0; i < ix.i64.size(); ++i) { + const int64_t src = ix.i64[i]; + if (src < 0 || src >= rows) fail(where() + ": index out of range"); + std::memcpy(r.data.data() + i * row, a.data.data() + size_t(src) * row, row * 4); + } + } else { + fail(where() + ": indices are not int64"); + } + } else if (n.op == "Equal") { + // Three forms in this family: the token window against <|endoftext|> (int64), and the causal-mask constant + // against itself (float, exact equality -- which is the ONNX semantic and the whole point of the pattern). + const Tensor& a = any_t(n.ins[0]); + const Tensor& b = any_t(n.ins[1]); + if (a.count() != b.count() && a.count() != 1 && b.count() != 1) + fail(where() + ": only scalar broadcast is implemented"); + r.dims = a.dims.size() >= b.dims.size() ? a.dims : b.dims; + const size_t n_out = std::max(a.count(), b.count()); + r.b.assign(n_out, 0); + for (size_t i = 0; i < n_out; ++i) { + const size_t ia = a.count() == 1 ? 0 : i, ib = b.count() == 1 ? 0 : i; + if (!a.i64.empty() && !b.i64.empty()) r.b[i] = a.i64[ia] == b.i64[ib] ? 1 : 0; + else if (!a.b.empty() && !b.b.empty()) r.b[i] = a.b[ia] == b.b[ib] ? 1 : 0; + else if (!a.data.empty() && !b.data.empty()) r.b[i] = a.data[ia] == b.data[ib] ? 1 : 0; + else fail(where() + ": operands are of different kinds"); + } + } else if (n.op == "Not") { + const Tensor& a = any_t(n.ins[0]); + if (a.b.empty()) fail(where() + ": operand is not a mask"); + r.dims = a.dims; + r.b.resize(a.b.size()); + for (size_t i = 0; i < a.b.size(); ++i) r.b[i] = a.b[i] ? 0 : 1; + } else if (n.op == "Where") { + // The causal-mask pattern in this family is Where(mask, -inf_scalar, 0_scalar): the *condition* carries the + // shape, so broadcasting only the two value operands would collapse the result to one element. + const Tensor& c = any_t(n.ins[0]); + const Tensor& a = any_t(n.ins[1]); + const Tensor& bb2 = any_t(n.ins[2]); + if (c.b.empty()) fail(where() + ": condition is not a mask"); + std::vector vals = plan_broadcast(a.dims, bb2.dims).dims; + r.dims = plan_broadcast(vals, c.dims).dims; + const size_t total = nelem(r.dims); + const bool use_i64 = !a.i64.empty(); + if (use_i64) r.i64.resize(total, 0); + else r.data.resize(total, 0.0f); + std::vector idx(r.dims.size(), 0); + for (size_t o = 0; o < total; ++o) { + const size_t ci = bcast_index(r.dims, c.dims, idx); + const size_t ia = bcast_index(r.dims, a.dims, idx); + const size_t ib = bcast_index(r.dims, bb2.dims, idx); + if (use_i64) r.i64[o] = c.b[ci] ? a.i64[ia] : bb2.i64[ib]; + else r.data[o] = c.b[ci] ? a.data[ia] : bb2.data[ib]; + for (size_t d = idx.size(); d-- > 0;) { + if (++idx[d] < size_t(r.dims[d])) break; + idx[d] = 0; + } + } + } else if (n.op == "Trilu") { + // The causal mask. ONNX zeroes the masked triangle, which for the float form is the "large constant" the + // following Equal/Where pair looks for; an int64 input is zeroed the same way. + const Tensor& a = any_t(n.ins[0]); + const int64_t k = n.attrs.count("k") ? n.attrs.at("k") : 0; + const bool upper = n.attrs.count("upper") ? n.attrs.at("upper") != 0 : true; + if (a.dims.size() < 2) fail(where() + ": needs a matrix"); + const int64_t R = a.dims[a.dims.size() - 2], C = a.dims.back(); + r.dims = a.dims; + r.data = a.data; + r.i64 = a.i64; + const size_t mat = size_t(R * C); + for (size_t base = 0; base + mat <= a.count(); base += mat) { + for (int64_t i = 0; i < R; ++i) { + for (int64_t j = 0; j < C; ++j) { + const bool keep = upper ? (j >= i + k) : (j <= i + k); + if (keep) continue; + if (!r.i64.empty()) r.i64[base + size_t(i * C + j)] = 0; + else r.data[base + size_t(i * C + j)] = 0.0f; + } + } + } + } else if (n.op == "Slice") { + const Tensor& a = any_t(n.ins[0]); + std::vector starts, ends, axes, steps; + auto shape_op = [&](size_t which, std::vector* dst) { + auto ii = inits.find(n.ins[which]); + if (ii == inits.end()) fail(where() + ": slice bounds must be initializers"); + load_init_i64(g, *ii->second, dst); + }; + if (n.ins.size() < 3) fail(where() + ": needs starts/ends"); + shape_op(1, &starts); + shape_op(2, &ends); + if (n.ins.size() > 3) shape_op(3, &axes); + if (n.ins.size() > 4) shape_op(4, &steps); + if (axes.empty()) + for (size_t i = 0; i < starts.size(); ++i) axes.push_back(int64_t(i)); + if (steps.empty()) + for (size_t i = 0; i < starts.size(); ++i) steps.push_back(1); + // Partial-rank slicing: an axis the node does not name is taken whole. Whisper's decoder slices one axis + // at a time (the position embedding, the logits' last position), so a full-rank-only rule refuses it. + std::vector from(a.dims.size(), 0), cnt = a.dims, step(a.dims.size(), 1); + for (size_t i = 0; i < axes.size(); ++i) { + const size_t ax = size_t(axes[i]); + if (ax >= a.dims.size()) fail(where() + ": axis out of range"); + int64_t st = starts[i] < 0 ? starts[i] + a.dims[ax] : starts[i]; + int64_t en = ends[i] < 0 ? ends[i] + a.dims[ax] : ends[i]; + st = std::max(0, std::min(st, a.dims[ax])); + en = std::max(st, std::min(en, a.dims[ax])); + from[ax] = st; + step[ax] = steps[i]; + cnt[ax] = (en - st + steps[i] - 1) / steps[i]; + } + for (size_t ax = 0; ax < step.size(); ++ax) + if (step[ax] != 1) fail(where() + ": only unit steps are implemented"); + r.dims = cnt; + r.data.resize(nelem(cnt)); + r.i64.resize(a.i64.empty() ? 0 : nelem(cnt)); + std::vector stride(a.dims.size(), 1); + for (size_t i = a.dims.size(); i-- > 1;) stride[i - 1] = stride[i] * size_t(a.dims[i]); + std::vector idx(cnt.size(), 0); + for (size_t o = 0; o < nelem(cnt); ++o) { + size_t src = 0; + for (size_t d = 0; d < idx.size(); ++d) src += size_t(from[d] + idx[d]) * stride[d]; + if (!a.i64.empty()) r.i64[o] = a.i64[src]; + else r.data[o] = a.data[src]; + for (size_t d = idx.size(); d-- > 0;) { + if (++idx[d] < cnt[d]) break; + idx[d] = 0; + } + } } else { fail(where() + " is not implemented; the pack needs an op this executor does not have"); } @@ -354,6 +517,8 @@ bool run_audio_graph(const OnnxGraphInfo& g, const float* mel, int64_t mel_bins, if (it == t.end()) fail("output '" + want + "' was never produced"); out->dims = it->second.dims; out->data = it->second.data; + out->i64 = it->second.i64; + out->b = it->second.b; return true; } catch (const std::exception& e) { if (why != nullptr) *why = e.what(); @@ -361,4 +526,45 @@ bool run_audio_graph(const OnnxGraphInfo& g, const float* mel, int64_t mel_bins, } } +} // namespace + +bool run_audio_graph(const OnnxGraphInfo& g, const float* mel, int64_t mel_bins, int64_t frames, const std::string& out_name, + AudioRunResult* out, std::string* why) { + std::map seed; + Tensor in; + in.dims = {1, mel_bins, frames}; + in.data.assign(mel, mel + size_t(mel_bins) * size_t(frames)); + if (g.inputs.empty()) { + if (why != nullptr) *why = "graph declares no input"; + return false; + } + seed[g.inputs[0]] = std::move(in); + return run_graph_impl(g, std::move(seed), out_name, out, why); +} + +bool run_decoder_graph(const OnnxGraphInfo& g, const int64_t* tokens, int64_t n_tokens, const float* xa, + const std::vector& xa_dims, const std::string& out_name, AudioRunResult* out, + std::string* why) { + if (g.inputs.size() < 1) { + if (why != nullptr) *why = "graph declares no input"; + return false; + } + std::map seed; + Tensor toks; + toks.dims = {1, n_tokens}; + toks.i64.assign(tokens, tokens + size_t(n_tokens)); + seed[g.inputs[0]] = std::move(toks); + if (g.inputs.size() > 1) { + if (xa == nullptr || xa_dims.empty()) { + if (why != nullptr) *why = "the graph has a second input (cross-attention) but none was supplied"; + return false; + } + Tensor x; + x.dims = xa_dims; + x.data.assign(xa, xa + nelem(xa_dims)); + seed[g.inputs[1]] = std::move(x); + } + return run_graph_impl(g, std::move(seed), out_name, out, why); +} + } // namespace onebit::npu::pack_import diff --git a/npu/audio_exec.h b/npu/audio_exec.h index 31ec1ee58..04433b4ef 100644 --- a/npu/audio_exec.h +++ b/npu/audio_exec.h @@ -19,14 +19,25 @@ namespace onebit::npu::pack_import { struct AudioRunResult { - std::vector dims; // dims of the requested output - std::vector data; // row-major + std::vector dims; // dims of the requested output + std::vector data; // row-major f32 payload + std::vector i64; // int64 payload, when the requested tensor is an index/token tensor + std::vector b; // bool payload, when it is a mask }; +// Two entry points, because a Whisper pack has two graphs: the encoder consumes the mel window as f32, the decoder consumes +// a token window as int64 plus the encoder's output. Everything else -- node order, weights, refusals -- is shared. +// // mel: (mel_bins, frames) row-major f32, the encoder's graph input in the layout its own first Conv consumes. // out_name: the graph output to return (empty = the graph's declared output). // Returns true and fills `out`; on failure returns false and sets `why` to the node and the reason. bool run_audio_graph(const OnnxGraphInfo& g, const float* mel, int64_t mel_bins, int64_t frames, const std::string& out_name, AudioRunResult* out, std::string* why); +// The decoder: `tokens` is the int64 token window for the graph's first input; `xa` (the encoder output, row-major f32) is +// bound to the graph's second input when the graph has one. `xa_dims` may be empty when there is no cross-attention input. +bool run_decoder_graph(const OnnxGraphInfo& g, const int64_t* tokens, int64_t n_tokens, const float* xa, + const std::vector& xa_dims, const std::string& out_name, AudioRunResult* out, + std::string* why); + } // namespace onebit::npu::pack_import diff --git a/tests/audio_exec_test.cpp b/tests/audio_exec_test.cpp index 4d5c163ef..1a63a6ee0 100644 --- a/tests/audio_exec_test.cpp +++ b/tests/audio_exec_test.cpp @@ -113,5 +113,50 @@ int main(int argc, char** argv) { return 1; } std::printf("audio exec: whole encoder on the host matches onnxruntime\n"); + + // ---- the decoder, on the same pack: 128-position token window x, the encoder output xa ---------------------------- + const char* dpath = (argc > 2) ? argv[2] : std::getenv("ONEBIT_TEST_WHISPER_DECODER"); + if (dpath == nullptr || *dpath == '\0') { + std::printf("SKIP decoder: pass the decoder .onnx as the second argument\n"); + return 0; + } + OnnxGraphInfo dg; + try { + dg = parse_onnx_graph(dpath); + } catch (const std::exception& e) { + std::printf("FAIL cannot parse %s: %s\n", dpath, e.what()); + return 1; + } + const int ntok = 128; + std::vector toks(size_t(ntok), 0); + toks[0] = 50258; // <|startoftranscript|> + toks[1] = 50359; // <|en|> + toks[2] = 50257; // <|transcribe|> + toks[3] = 1234; + std::vector xad; + for (int64_t d : r.dims) xad.push_back(d); + AudioRunResult dr; + if (!run_decoder_graph(dg, toks.data(), ntok, r.data.data(), xad, "", &dr, &why)) { + std::printf("FAIL decoder run: %s\n", why.c_str()); + return 1; + } + std::printf("decoder ran %zu nodes -> dims", dg.nodes.size()); + for (int64_t d : dr.dims) std::printf(" %lld", (long long)d); + std::printf("\n"); + check(dr.dims.size() == 3 && dr.dims[0] == 1 && dr.dims[1] == ntok, "decoder output is [1,128,vocab]"); + if (dr.dims.size() == 3) { + const int64_t V = dr.dims[2]; + auto dl = [&](int64_t t, int64_t v) { return dr.data[size_t(t) * size_t(V) + size_t(v)]; }; + // onnxruntime, same chain (formula mel -> encoder -> xa -> decoder): the decoder's last row and row 3. + const double last[5] = {8.26411438, -1.29840982, -2.36972690, -3.26691365, -2.12733102}; + const double row3[4] = {18.45385361, 17.14962196, 12.60445786, 10.96115971}; + for (int i = 0; i < 5; ++i) expect_close(dl(ntok - 1, i), last[i], 2e-3, "decoder last row col " + std::to_string(i)); + for (int i = 0; i < 4; ++i) expect_close(dl(3, i), row3[i], 2e-3, "decoder row 3 col " + std::to_string(i)); + } + if (g_failures != 0) { + std::printf("audio exec: %d failures\n", g_failures); + return 1; + } + std::printf("audio exec: decoder graph runs on the host and matches onnxruntime\n"); return 0; } From 351123ff57af53d44ffc92067915260416aacee9 Mon Sep 17 00:00:00 2001 From: bong-water-water-bong Date: Tue, 10 Nov 2026 11:20:55 -0400 Subject: [PATCH 16/20] npu: the log-mel front-end, checked against the reference extractor (engine#1931) --- CMakeLists.txt | 3 + npu/forward/include/audio_mel.h | 126 ++++++++++++++++++++++++++++++++ tests/audio_mel_test.cpp | 97 ++++++++++++++++++++++++ 3 files changed, 226 insertions(+) create mode 100644 npu/forward/include/audio_mel.h create mode 100644 tests/audio_mel_test.cpp diff --git a/CMakeLists.txt b/CMakeLists.txt index dfd0c1c7d..f5fbe9af5 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -1084,6 +1084,8 @@ target_include_directories(audio_ops_test PRIVATE npu/forward/include) # census the host reference was written against. Skips (77) without a path. add_executable(onnx_graph_test tests/onnx_graph_test.cpp) target_link_libraries(onnx_graph_test PRIVATE onebit_npu) +add_executable(audio_mel_test tests/audio_mel_test.cpp) +target_include_directories(audio_mel_test PRIVATE npu/forward/include) add_executable(audio_exec_test tests/audio_exec_test.cpp) target_link_libraries(audio_exec_test PRIVATE onebit_npu) add_executable(flm_head_import tools/flm_head_import.cpp) @@ -1310,6 +1312,7 @@ if(ONEBIT_NPU_MODEL AND ONEBIT_NPU_BO_DUMP) add_test(NAME npu_pack_model COMMAND npu_pack_test ${ONEBIT_NPU_MODEL} ${ONEBIT_NPU_BO_DUMP}) endif() add_test(NAME audio_ops COMMAND audio_ops_test) +add_test(NAME audio_mel COMMAND audio_mel_test) set_tests_properties(audio_ops PROPERTIES SKIP_RETURN_CODE 77) add_test(NAME onnx_graph COMMAND onnx_graph_test ${ONEBIT_AUDIO_ONNX}) add_test(NAME audio_exec COMMAND audio_exec_test ${ONEBIT_AUDIO_ONNX}) diff --git a/npu/forward/include/audio_mel.h b/npu/forward/include/audio_mel.h new file mode 100644 index 000000000..2f7919891 --- /dev/null +++ b/npu/forward/include/audio_mel.h @@ -0,0 +1,126 @@ +// Copyright 2026 bong-water-water-bong +// SPDX-License-Identifier: Apache-2.0 +// +// audio_mel.h -- the log-mel front-end: hann-windowed STFT (n_fft 400, hop 160, reflect padding), power spectrum, a +// slaney-normalised mel filterbank, then log10 and the (max - 8, +4)/4 scaling a Whisper encoder's input expects. +// +// This is the one part of the pipeline that is *not* in the pack: the pack's graph starts at the mel window, so the numbers +// here have to match the reference extractor's, and tests/audio_mel_test.cpp checks them against it. +// +// The filterbank is built, not shipped: librosa's slaney mel scale (f_sp = 200/3 Hz, linear below 1 kHz, log above, +// logstep = ln(6.4)/27) with slaney normalisation, which is what Whisper's feature extractor uses. + +#pragma once + +#include +#include +#include + +namespace onebit::audio { + +inline double hz_to_mel(double hz) { + const double f_sp = 200.0 / 3.0; + const double min_log_hz = 1000.0, min_log_mel = min_log_hz / f_sp, logstep = std::log(6.4) / 27.0; + if (hz < min_log_hz) return hz / f_sp; + return min_log_mel + std::log(hz / min_log_hz) / logstep; +} + +inline double mel_to_hz(double mel) { + const double f_sp = 200.0 / 3.0; + const double min_log_hz = 1000.0, min_log_mel = min_log_hz / f_sp, logstep = std::log(6.4) / 27.0; + if (mel < min_log_mel) return mel * f_sp; + return min_log_hz * std::exp(logstep * (mel - min_log_mel)); +} + +// (bins, 1 + n_fft/2): triangular filters, slaney-normalised, spanning 0..nyquist. +inline std::vector mel_filterbank(int bins, int n_fft, int sample_rate) { + const int nfreqs = n_fft / 2 + 1; + const double nyq = sample_rate / 2.0; + const double mlo = hz_to_mel(0.0), mhi = hz_to_mel(nyq); + std::vector f(bins + 2); + for (int i = 0; i < bins + 2; ++i) f[i] = mel_to_hz(mlo + (mhi - mlo) * double(i) / double(bins + 1)); + // FFT bin centres in Hz, and the per-filter bandwidth normalisation. + std::vector hz(nfreqs); + for (int k = 0; k < nfreqs; ++k) hz[k] = double(k) * nyq / double(nfreqs - 1); + std::vector enorm(bins); + for (int i = 0; i < bins; ++i) { + const double w = f[i + 2] - f[i]; + enorm[i] = w > 0 ? 2.0 / w : 0.0; + } + std::vector m(size_t(bins) * size_t(nfreqs), 0.0f); + for (int i = 0; i < bins; ++i) { + const double lo = f[i], ctr = f[i + 1], hi = f[i + 2]; + for (int k = 0; k < nfreqs; ++k) { + const double x = hz[k]; + double v = 0.0; + if (x >= lo && x <= ctr && ctr > lo) v = (x - lo) / (ctr - lo); + else if (x > ctr && x <= hi && hi > ctr) v = (hi - x) / (hi - ctr); + m[size_t(i) * nfreqs + k] = float(v * enorm[i]); + } + } + return m; +} + +struct MelConfig { + int sample_rate = 16000; + int n_fft = 400; + int hop = 160; + int bins = 80; + int frames = 3000; // the encoder's window: 30 s +}; + +// audio: `samples` mono f32 in [-1,1]. out: (bins, frames) row-major. Returns false only on a size the config cannot fill. +inline bool log_mel_spectrogram(const float* audio, size_t samples, const MelConfig& cfg, std::vector* out) { + const int pad = cfg.n_fft / 2; + const int nfreqs = cfg.n_fft / 2 + 1; + if (samples + 2 * size_t(pad) < size_t(cfg.n_fft)) return false; + + // reflect padding, both sides (torch's center=True default for stft) + std::vector x(samples + 2 * size_t(pad)); + for (int i = 0; i < pad; ++i) x[size_t(i)] = audio[pad - i]; + for (size_t i = 0; i < samples; ++i) x[size_t(pad) + i] = audio[i]; + for (int i = 0; i < pad; ++i) x[size_t(pad) + samples + size_t(i)] = audio[samples - 2 - size_t(i)]; + + std::vector win(cfg.n_fft); + for (int i = 0; i < cfg.n_fft; ++i) // periodic hann, as torch.hann_window(w) is by default + win[size_t(i)] = 0.5 - 0.5 * std::cos(2.0 * M_PI * double(i) / double(cfg.n_fft)); + + const std::vector fb = mel_filterbank(cfg.bins, cfg.n_fft, cfg.sample_rate); + out->assign(size_t(cfg.bins) * size_t(cfg.frames), 0.0f); + std::vector re(nfreqs), im(nfreqs), power(nfreqs); + for (int t = 0; t < cfg.frames; ++t) { + const size_t base = size_t(t) * size_t(cfg.hop); + for (int k = 0; k < nfreqs; ++k) { + double sr = 0.0, si = 0.0; + for (int n = 0; n < cfg.n_fft; ++n) { + const double v = x[base + size_t(n)] * win[size_t(n)]; + const double ang = -2.0 * M_PI * double(k) * double(n) / double(cfg.n_fft); + sr += v * std::cos(ang); + si += v * std::sin(ang); + } + re[size_t(k)] = sr; + im[size_t(k)] = si; + power[size_t(k)] = sr * sr + si * si; + } + for (int b = 0; b < cfg.bins; ++b) { + const float* row = fb.data() + size_t(b) * nfreqs; + double acc = 0.0; + for (int k = 0; k < nfreqs; ++k) acc += double(row[k]) * power[size_t(k)]; + out->at(size_t(b) * size_t(cfg.frames) + size_t(t)) = float(acc); + } + } + // log10 with a floor, then Whisper's dynamic-range clamp and scale. + double mx = -1e30; + for (float& v : *out) { + v = float(std::log10(std::max(double(v), 1e-10))); + if (v > mx) mx = v; + } + const double floor_db = mx - 8.0; + for (float& v : *out) { + const double c = std::max(double(v), floor_db); + v = float((c + 4.0) / 4.0); + } + return true; +} + +} // namespace onebit::audio diff --git a/tests/audio_mel_test.cpp b/tests/audio_mel_test.cpp new file mode 100644 index 000000000..70d7e823b --- /dev/null +++ b/tests/audio_mel_test.cpp @@ -0,0 +1,97 @@ +// Copyright 2026 bong-water-water-bong +// SPDX-License-Identifier: Apache-2.0 +// +// audio_mel_test.cpp -- the log-mel front-end against the reference extractor's numbers. +// +// Input: 30 s at 16 kHz of three summed sinusoids (440, 1250, 3100 Hz), built from a formula so both sides use the same +// signal. Expected values are transformers' WhisperFeatureExtractor (feature_size=80) on that signal. + +#include +#include +#include +#include + +#include "audio_mel.h" + +namespace { + +int g_failures = 0; + +void check(bool cond, const std::string& what) { + if (!cond) { + std::printf("FAIL %s\n", what.c_str()); + ++g_failures; + } +} + +void expect_close(float got, double want, double tol, const std::string& what) { + const double d = std::fabs(double(got) - want); + if (!(d <= tol)) { + std::printf("FAIL %s: got %.9g want %.9g (|d| %.3g)\n", what.c_str(), double(got), want, d); + ++g_failures; + } +} + +} // namespace + +int main() { + onebit::audio::MelConfig cfg; + const int sr = cfg.sample_rate; + const size_t n = size_t(sr) * 30; + std::vector sig(n); + for (size_t i = 0; i < n; ++i) { + const double t = double(i) / sr; + sig[i] = float(0.3 * std::sin(2 * M_PI * 440.0 * t) + 0.2 * std::sin(2 * M_PI * 1250.0 * t) + + 0.1 * std::sin(2 * M_PI * 3100.0 * t)); + } + + std::vector mel; + if (!onebit::audio::log_mel_spectrogram(sig.data(), sig.size(), cfg, &mel)) { + std::printf("FAIL front-end refused a 30 s window\n"); + return 1; + } + check(mel.size() == size_t(cfg.bins) * size_t(cfg.frames), "mel is (80, 3000)"); + if (mel.size() != size_t(cfg.bins) * size_t(cfg.frames)) return 1; + + double mean = 0.0, mn = 1e30, mx = -1e30; + for (float v : mel) { + mean += v; + mn = std::min(mn, v); + mx = std::max(mx, v); + } + mean /= double(mel.size()); + auto at = [&](int b, int t) { return mel[size_t(b) * size_t(cfg.frames) + size_t(t)]; }; + + // reference: mean -0.33795589, min -0.672718, max 1.327282 + expect_close(float(mean), -0.33795589, 2e-4, "mel mean"); + expect_close(float(mn), -0.672718, 2e-4, "mel min"); + expect_close(float(mx), 1.327282, 2e-4, "mel max"); + + // bin 0, first five frames -- the onset and the dynamic-range clamp (three equal -0.6727176 values) + const double p0[5] = {0.92388016, 0.41213650, -0.67271757, -0.67271757, -0.67271757}; + for (int i = 0; i < 5; ++i) expect_close(at(0, i), p0[i], 2e-3, "bin 0 frame " + std::to_string(i)); + // mid band, the end of the window + const double p1[5] = {-0.65874779, -0.66997766, -0.65873981, -0.66998839, -0.65872025}; + for (int i = 0; i < 5; ++i) expect_close(at(40, 1499 + i), p1[i], 2e-3, "bin 40 frame " + std::to_string(1499 + i)); + // top bin + const double p2[3] = {0.18210298, -0.32799470, -0.67271757}; + for (int i = 0; i < 3; ++i) expect_close(at(79, i), p2[i], 2e-3, "bin 79 frame " + std::to_string(i)); + + // A silent window still produces a full, finite spectrogram: the log floor and the max-8 clamp own that case, and a + // NaN there would poison every frame the encoder sees. + std::vector silent(n, 0.0f), smel; + check(onebit::audio::log_mel_spectrogram(silent.data(), silent.size(), cfg, &smel), "silence is accepted"); + bool finite = true; + for (float v : smel) finite = finite && std::isfinite(v); + check(finite, "a silent window produces finite numbers"); + // Silence is the floor of the whole chain: power 0 -> log10(1e-10) = -10, the max is that too, so the max-8 clamp never + // bites and the scale gives (-10+4)/4 = -1.5. (My first expectation here said 1.0 -- the test caught its author again.) + expect_close(smel[0], -1.5f, 1e-6, "silence lands on the log floor: (-10+4)/4 == -1.5"); + + if (g_failures != 0) { + std::printf("audio mel: %d failures\n", g_failures); + return 1; + } + std::printf("audio mel: front-end matches the reference extractor on tones, silences on the clamp\n"); + return 0; +} From 324d57d220a89359e2645185593b9f416d26e859 Mon Sep 17 00:00:00 2001 From: bong-water-water-bong Date: Tue, 10 Nov 2026 11:44:47 -0400 Subject: [PATCH 17/20] npu: whisper_transcribe -- wav to token ids on the host, matching the reference decode (engine#1931) --- CMakeLists.txt | 3 + tools/whisper_transcribe.cpp | 174 +++++++++++++++++++++++++++++++++++ 2 files changed, 177 insertions(+) create mode 100644 tools/whisper_transcribe.cpp diff --git a/CMakeLists.txt b/CMakeLists.txt index f5fbe9af5..dfd7ff673 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -1084,6 +1084,9 @@ target_include_directories(audio_ops_test PRIVATE npu/forward/include) # census the host reference was written against. Skips (77) without a path. add_executable(onnx_graph_test tests/onnx_graph_test.cpp) target_link_libraries(onnx_graph_test PRIVATE onebit_npu) +add_executable(whisper_transcribe tools/whisper_transcribe.cpp) +target_include_directories(whisper_transcribe PRIVATE npu/forward/include) +target_link_libraries(whisper_transcribe PRIVATE onebit_npu) add_executable(audio_mel_test tests/audio_mel_test.cpp) target_include_directories(audio_mel_test PRIVATE npu/forward/include) add_executable(audio_exec_test tests/audio_exec_test.cpp) diff --git a/tools/whisper_transcribe.cpp b/tools/whisper_transcribe.cpp new file mode 100644 index 000000000..647a68398 --- /dev/null +++ b/tools/whisper_transcribe.cpp @@ -0,0 +1,174 @@ +// Copyright 2026 bong-water-water-bong +// SPDX-License-Identifier: Apache-2.0 +// +// whisper_transcribe.cpp -- the whole host pipeline: wav -> log-mel -> encoder -> greedy decode -> token ids. +// +// ./whisper_transcribe [expected-ids-csv ...] +// +// A tool rather than a test because one forward pass of each graph is ~26 s at these shapes and a sentence is 25 of them +// (see PR #1931); `ctest` would time out for a reason that has nothing to do with correctness. The numbers it prints are +// the ones a test would assert, so `expected-ids-csv` turns it into a check when run by hand. +// +// The pack on disk carries no tokenizer, so this stops at ids: `400 370 452 ...` is the sentence, and comparing ids is a +// stricter check than comparing text (two ids can share a detokenisation, and neither can share an id). + +#include +#include +#include +#include +#include +#include + +#include "audio_exec.h" +#include "audio_mel.h" +#include "onnx_graph.h" + +namespace { + +std::vector read_wav_16k_mono(const char* path, std::string* why) { + FILE* f = std::fopen(path, "rb"); + if (f == nullptr) { + *why = std::string("cannot open ") + path; + return {}; + } + std::vector raw; + uint8_t buf[65536]; + size_t n; + while ((n = std::fread(buf, 1, sizeof buf, f)) > 0) raw.insert(raw.end(), buf, buf + n); + std::fclose(f); + if (raw.size() < 44 || std::memcmp(raw.data(), "RIFF", 4) != 0 || std::memcmp(raw.data() + 8, "WAVE", 4) != 0) { + *why = "not a RIFF/WAVE file"; + return {}; + } + size_t pos = 12; + uint16_t channels = 0, bits = 0; + uint32_t rate = 0; + const uint8_t* pcm = nullptr; + size_t pcm_len = 0; + while (pos + 8 <= raw.size()) { + const uint32_t len = uint32_t(raw[pos + 4]) | (uint32_t(raw[pos + 5]) << 8) | (uint32_t(raw[pos + 6]) << 16) | + (uint32_t(raw[pos + 7]) << 24); + const size_t body = pos + 8; + if (std::memcmp(raw.data() + pos, "fmt ", 4) == 0 && body + 16 <= raw.size()) { + channels = uint16_t(raw[body + 2] | (raw[body + 3] << 8)); + rate = uint32_t(raw[body + 4]) | (uint32_t(raw[body + 5]) << 8) | (uint32_t(raw[body + 6]) << 16) | + (uint32_t(raw[body + 7]) << 24); + bits = uint16_t(raw[body + 14] | (raw[body + 15] << 8)); + } else if (std::memcmp(raw.data() + pos, "data", 4) == 0) { + pcm = raw.data() + body; + pcm_len = std::min(len, raw.size() - body); + } + pos = body + len + (len & 1); + } + if (pcm == nullptr || channels == 0) { + *why = "no data chunk"; + return {}; + } + if (rate != 16000) { + *why = "expected 16 kHz, got " + std::to_string(rate); + return {}; + } + if (bits != 16) { + *why = "expected 16-bit PCM, got " + std::to_string(bits) + "-bit"; + return {}; + } + const size_t frames = pcm_len / (2 * size_t(channels)); + std::vector out(frames); + for (size_t i = 0; i < frames; ++i) { + int32_t acc = 0; + for (uint16_t c = 0; c < channels; ++c) { + const uint8_t* s = pcm + (i * channels + c) * 2; + acc += int16_t(uint16_t(s[0] | (s[1] << 8))); + } + out[i] = float(double(acc) / double(channels) / 32768.0); + } + return out; +} + +} // namespace + +int main(int argc, char** argv) { + if (argc < 4) { + std::fprintf(stderr, "usage: %s [expected-id ...]\n", argv[0]); + return 2; + } + using namespace onebit::npu::pack_import; + + std::string why; + std::vector audio = read_wav_16k_mono(argv[3], &why); + if (audio.empty()) { + std::fprintf(stderr, "wav: %s\n", why.c_str()); + return 1; + } + // Whisper's window is a fixed 30 s: pad with silence, do not scale. + onebit::audio::MelConfig cfg; + const size_t want = size_t(cfg.sample_rate) * size_t(cfg.frames) * size_t(cfg.hop) / size_t(cfg.frames); + audio.resize(std::max(audio.size(), want), 0.0f); + + std::vector mel; + if (!onebit::audio::log_mel_spectrogram(audio.data(), audio.size(), cfg, &mel)) { + std::fprintf(stderr, "mel: refused\n"); + return 1; + } + OnnxGraphInfo enc, dec; + try { + enc = parse_onnx_graph(argv[1]); + dec = parse_onnx_graph(argv[2]); + } catch (const std::exception& e) { + std::fprintf(stderr, "graph: %s\n", e.what()); + return 1; + } + AudioRunResult xa; + if (!run_audio_graph(enc, mel.data(), cfg.bins, cfg.frames, "", &xa, &why)) { + std::fprintf(stderr, "encoder: %s\n", why.c_str()); + return 1; + } + std::printf("encoder: ["); + for (int64_t d : xa.dims) std::printf(" %lld", (long long)d); + std::printf(" ]\n"); + + // <|startoftranscript|>, <|en|>, <|transcribe|>, <|notimestamps|> -- then greedy, reading the logits of the last real + // prompt position (HF pads the 448/128-token window with <|endoftext|>, so the prompt is a prefix, not a suffix). + const int64_t kSot = 50258, kEn = 50259, kTranscribe = 50359, kNoTimestamps = 50363, kEot = 50257; + const int window = 128; + std::vector seq = {kSot, kEn, kTranscribe, kNoTimestamps}; + std::vector ids; + for (int step = 0; step < 60; ++step) { + std::vector toks(size_t(window), kEot); + for (size_t i = 0; i < seq.size() && i < size_t(window); ++i) toks[i] = seq[i]; + AudioRunResult y; + if (!run_decoder_graph(dec, toks.data(), window, xa.data.data(), xa.dims, "", &y, &why)) { + std::fprintf(stderr, "decoder at step %d: %s\n", step, why.c_str()); + return 1; + } + const int64_t vocab = y.dims.empty() ? 0 : y.dims.back(); + const size_t row = (seq.size() - 1) * size_t(vocab); + int64_t best = 0; + for (int64_t v = 1; v < vocab; ++v) + if (y.data[row + size_t(v)] > y.data[row + size_t(best)]) best = v; + if (best == kEot || seq.size() >= size_t(window)) break; + ids.push_back(best); + seq.push_back(best); + } + std::printf("ids:"); + for (int64_t v : ids) std::printf(" %lld", (long long)v); + std::printf("\n"); + + if (argc > 4) { + int bad = 0; + const int n = argc - 4; + if (int(ids.size()) != n) { + std::printf("FAIL length: got %zu want %d\n", ids.size(), n); + ++bad; + } else { + for (int i = 0; i < n; ++i) + if (ids[size_t(i)] != std::atoll(argv[4 + i])) { + std::printf("FAIL id[%d]: got %lld want %s\n", i, (long long)ids[size_t(i)], argv[4 + i]); + ++bad; + } + } + if (bad == 0) std::printf("transcript ids match the reference decode\n"); + return bad == 0 ? 0 : 1; + } + return 0; +} From 56c35aa3b71527da364ff6572d2544d3eac51d24 Mon Sep 17 00:00:00 2001 From: bong-water-water-bong Date: Tue, 10 Nov 2026 13:19:13 -0400 Subject: [PATCH 18/20] npu: whisper_transcribe reports where the host time goes (engine#1931) --- tools/whisper_transcribe.cpp | 11 +++++++++++ 1 file changed, 11 insertions(+) diff --git a/tools/whisper_transcribe.cpp b/tools/whisper_transcribe.cpp index 647a68398..004c74d42 100644 --- a/tools/whisper_transcribe.cpp +++ b/tools/whisper_transcribe.cpp @@ -12,6 +12,7 @@ // The pack on disk carries no tokenizer, so this stops at ids: `400 370 452 ...` is the sentence, and comparing ids is a // stricter check than comparing text (two ids can share a detokenisation, and neither can share an id). +#include #include #include #include @@ -105,11 +106,13 @@ int main(int argc, char** argv) { const size_t want = size_t(cfg.sample_rate) * size_t(cfg.frames) * size_t(cfg.hop) / size_t(cfg.frames); audio.resize(std::max(audio.size(), want), 0.0f); + const auto t_start = std::chrono::steady_clock::now(); std::vector mel; if (!onebit::audio::log_mel_spectrogram(audio.data(), audio.size(), cfg, &mel)) { std::fprintf(stderr, "mel: refused\n"); return 1; } + const auto t_mel = std::chrono::steady_clock::now(); OnnxGraphInfo enc, dec; try { enc = parse_onnx_graph(argv[1]); @@ -123,6 +126,7 @@ int main(int argc, char** argv) { std::fprintf(stderr, "encoder: %s\n", why.c_str()); return 1; } + const auto t_enc = std::chrono::steady_clock::now(); std::printf("encoder: ["); for (int64_t d : xa.dims) std::printf(" %lld", (long long)d); std::printf(" ]\n"); @@ -150,6 +154,13 @@ int main(int argc, char** argv) { ids.push_back(best); seq.push_back(best); } + const auto t_dec = std::chrono::steady_clock::now(); + const auto secs = [](std::chrono::steady_clock::time_point a, std::chrono::steady_clock::time_point b) { + return std::chrono::duration(b - a).count(); + }; + std::printf("time: mel+load %.3f s, encoder %.3f s, decode %.3f s for %zu token(s) -> %.3f s/token, wall %.3f s\n", + secs(t_start, t_mel), secs(t_mel, t_enc), secs(t_enc, t_dec), ids.size(), + ids.empty() ? 0.0 : secs(t_enc, t_dec) / double(ids.size()), secs(t_start, t_dec)); std::printf("ids:"); for (int64_t v : ids) std::printf(" %lld", (long long)v); std::printf("\n"); From 61426f5c6c6292ea4094f1b1c86e474f7cf4acdb Mon Sep 17 00:00:00 2001 From: bong-water-water-bong Date: Tue, 10 Nov 2026 13:56:49 -0400 Subject: [PATCH 19/20] npu: split independent MatMul rows across threads, bit-identical, measured at last (engine#1931) --- npu/audio_exec.cpp | 64 ++++++++++++++++++++++++++++++++++------------ 1 file changed, 48 insertions(+), 16 deletions(-) diff --git a/npu/audio_exec.cpp b/npu/audio_exec.cpp index a2b73c7e5..abeab9eb1 100644 --- a/npu/audio_exec.cpp +++ b/npu/audio_exec.cpp @@ -9,7 +9,9 @@ #include #include #include +#include #include +#include #include #include "forward/include/audio_ops.h" @@ -80,6 +82,33 @@ size_t bcast_index(const std::vector& od, const std::vector& s return off; } +// Deterministic fork-join over independent units of work (output rows, batch indices). Each unit writes only its own slice, so +// every element's arithmetic is unchanged by the thread count: the result is bit-identical to the serial loop, which is the +// property a reference implementation has to keep. ONEBIT_AUDIO_THREADS=1 forces the serial path (the control for measuring it). +template +void parallel_units(size_t units, Fn&& fn) { + size_t want = 1; + if (const char* env = std::getenv("ONEBIT_AUDIO_THREADS")) { + const long n = std::strtol(env, nullptr, 10); + want = n > 0 ? size_t(n) : 1; + } else { + const size_t hw = std::thread::hardware_concurrency(); + want = hw == 0 ? 1 : (hw > 8 ? 8 : hw); + } + if (want > units) want = units; + if (want <= 1) { + for (size_t u = 0; u < units; ++u) fn(u); + return; + } + std::vector pool; + pool.reserve(want); + for (size_t t = 0; t < want; ++t) + pool.emplace_back([&, t]() { + for (size_t u = t; u < units; u += want) fn(u); + }); + for (auto& th : pool) th.join(); +} + void broadcast_binary(const Tensor& a, const Tensor& b, const std::string& op, Tensor* out) { const Broadcast p = plan_broadcast(a.dims, b.dims); out->dims = p.dims; @@ -229,9 +258,19 @@ bool run_graph_impl(const OnnxGraphInfo& g, std::map t, con r.dims.push_back(M); r.dims.push_back(N); r.data.assign(nelem(r.dims), 0.0f); - std::vector idx(bp.dims.size(), 0); + // Rows are independent, so the work splits over (batch, row) without touching any element's arithmetic. const size_t batches = nelem(bp.dims); - for (size_t bi = 0; bi < batches; ++bi) { + const size_t rows = size_t(std::max(M, 1)); + parallel_units(batches * rows, [&](size_t u) { + const size_t bi = u / rows; + const int64_t i = int64_t(u % rows); + std::vector idx(bp.dims.size(), 0); + size_t rem = bi; + for (size_t d = bp.dims.size(); d-- > 0;) { // row-major unflatten of the batch index + const size_t dim = size_t(bp.dims[d] ? bp.dims[d] : 1); + idx[d] = rem % dim; + rem /= dim; + } size_t oa = 0, ob = 0; for (size_t d = 0; d < idx.size(); ++d) { oa += idx[d] * bp.stride_big[d]; @@ -239,21 +278,14 @@ bool run_graph_impl(const OnnxGraphInfo& g, std::map t, con } const float* ap = a.data.data() + oa * size_t(M * K); const float* bpp = b.data.data() + ob * size_t(K * N); - float* rp = r.data.data() + bi * size_t(M * N); - for (int64_t i = 0; i < M; ++i) { - float* rr = rp + size_t(i) * size_t(N); - for (int64_t k = 0; k < K; ++k) { - const float av = ap[size_t(i) * size_t(K) + size_t(k)]; - if (av == 0.0f) continue; - const float* br = bpp + size_t(k) * size_t(N); - for (int64_t j = 0; j < N; ++j) rr[j] += av * br[j]; - } - } - for (size_t d = idx.size(); d-- > 0;) { - if (++idx[d] < size_t(bp.dims[d])) break; - idx[d] = 0; + float* rr = r.data.data() + bi * size_t(M * N) + size_t(i) * size_t(N); + for (int64_t k = 0; k < K; ++k) { + const float av = ap[size_t(i) * size_t(K) + size_t(k)]; + if (av == 0.0f) continue; + const float* br = bpp + size_t(k) * size_t(N); + for (int64_t j = 0; j < N; ++j) rr[j] += av * br[j]; } - } + }); } else if (n.op == "Add" || n.op == "Mul" || n.op == "Div" || n.op == "Sub") { // An operand is either a tensor an earlier node produced or an initializer (a bias vector, the 1/sqrt(d) // scalar); init_f32() handles the second without a separate code path at every call site. From 49b27b4019ae3c28f7eb071415a0aff4f2963c3d Mon Sep 17 00:00:00 2001 From: bong-water-water-bong Date: Tue, 10 Nov 2026 17:56:12 -0400 Subject: [PATCH 20/20] npu: the audio fork-join is tested for bit-identity and survives a throwing unit; honest timing labels; notices (engine#1931) Answers the three open reviews on #1931 at 61426f5: - parallel_units moves to npu/audio_parallel.h; a worker that throws is caught, every worker is joined, the lowest worker's exception is rethrown. - tests/audio_parallel_test.cpp: the executor's MatMul compared with memcmp at 1/2/3/7/8/64 threads and unset, on an in-memory graph (no pack); every unit runs exactly once; a throwing unit is a caught exception. - whisper_transcribe reports wav, mel, load, encoder, decode separately: 'mel+load' did not include the wav read, and 'encoder' included both graph parses. - The audio tests are build dependencies of onebit, so CI builds them without a workflow edit; full Apache notice on the ten audio files. Co-Authored-By: Claude Fable 5.2 --- CMakeLists.txt | 8 ++ npu/audio_exec.cpp | 41 +++---- npu/audio_exec.h | 12 ++ npu/audio_parallel.h | 81 +++++++++++++ npu/forward/include/audio_mel.h | 12 ++ npu/forward/include/audio_ops.h | 12 ++ npu/onnx_graph.h | 12 ++ tests/audio_exec_test.cpp | 12 ++ tests/audio_mel_test.cpp | 12 ++ tests/audio_ops_test.cpp | 12 ++ tests/audio_parallel_test.cpp | 198 ++++++++++++++++++++++++++++++++ tests/onnx_graph_test.cpp | 12 ++ tools/whisper_transcribe.cpp | 24 +++- 13 files changed, 417 insertions(+), 31 deletions(-) create mode 100644 npu/audio_parallel.h create mode 100644 tests/audio_parallel_test.cpp diff --git a/CMakeLists.txt b/CMakeLists.txt index dab6ec497..2f00c3af6 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -1223,6 +1223,13 @@ add_executable(audio_mel_test tests/audio_mel_test.cpp) target_include_directories(audio_mel_test PRIVATE npu/forward/include) add_executable(audio_exec_test tests/audio_exec_test.cpp) target_link_libraries(audio_exec_test PRIVATE onebit_npu) +# The executor's fork-join (npu/audio_parallel.h): MatMul bit-identical across ONEBIT_AUDIO_THREADS, and a throwing unit +# is a caught exception rather than std::terminate. The graph is built in memory, so: no pack, no device. +add_executable(audio_parallel_test tests/audio_parallel_test.cpp) +target_link_libraries(audio_parallel_test PRIVATE onebit_npu Threads::Threads) +# ci.yml's build list is not ours to edit without `workflow` scope, so these reach CI the way ops_dump_test does: as build +# dependencies of a target it already builds (tools/ci/ci_targets_check.py follows this edge). +add_dependencies(onebit audio_ops_test audio_mel_test onnx_graph_test audio_exec_test audio_parallel_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) @@ -1500,6 +1507,7 @@ add_test(NAME audio_mel COMMAND audio_mel_test) set_tests_properties(audio_ops PROPERTIES SKIP_RETURN_CODE 77) add_test(NAME onnx_graph COMMAND onnx_graph_test ${ONEBIT_AUDIO_ONNX}) add_test(NAME audio_exec COMMAND audio_exec_test ${ONEBIT_AUDIO_ONNX}) +add_test(NAME audio_parallel COMMAND audio_parallel_test) set_tests_properties(audio_exec PROPERTIES SKIP_RETURN_CODE 77) set_tests_properties(onnx_graph PROPERTIES SKIP_RETURN_CODE 77) # The private NPU route registry (npu/private_route.h), without an add-on. diff --git a/npu/audio_exec.cpp b/npu/audio_exec.cpp index abeab9eb1..edb394ad7 100644 --- a/npu/audio_exec.cpp +++ b/npu/audio_exec.cpp @@ -1,9 +1,22 @@ // 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. +// // audio_exec.cpp -- the host f32 executor for an audio pack's graph (see audio_exec.h for the contract). #include "audio_exec.h" +#include "audio_parallel.h" #include #include @@ -11,7 +24,6 @@ #include #include #include -#include #include #include "forward/include/audio_ops.h" @@ -82,33 +94,6 @@ size_t bcast_index(const std::vector& od, const std::vector& s return off; } -// Deterministic fork-join over independent units of work (output rows, batch indices). Each unit writes only its own slice, so -// every element's arithmetic is unchanged by the thread count: the result is bit-identical to the serial loop, which is the -// property a reference implementation has to keep. ONEBIT_AUDIO_THREADS=1 forces the serial path (the control for measuring it). -template -void parallel_units(size_t units, Fn&& fn) { - size_t want = 1; - if (const char* env = std::getenv("ONEBIT_AUDIO_THREADS")) { - const long n = std::strtol(env, nullptr, 10); - want = n > 0 ? size_t(n) : 1; - } else { - const size_t hw = std::thread::hardware_concurrency(); - want = hw == 0 ? 1 : (hw > 8 ? 8 : hw); - } - if (want > units) want = units; - if (want <= 1) { - for (size_t u = 0; u < units; ++u) fn(u); - return; - } - std::vector pool; - pool.reserve(want); - for (size_t t = 0; t < want; ++t) - pool.emplace_back([&, t]() { - for (size_t u = t; u < units; u += want) fn(u); - }); - for (auto& th : pool) th.join(); -} - void broadcast_binary(const Tensor& a, const Tensor& b, const std::string& op, Tensor* out) { const Broadcast p = plan_broadcast(a.dims, b.dims); out->dims = p.dims; diff --git a/npu/audio_exec.h b/npu/audio_exec.h index 04433b4ef..df7e099ce 100644 --- a/npu/audio_exec.h +++ b/npu/audio_exec.h @@ -1,6 +1,18 @@ // 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. +// // audio_exec.h -- run an audio pack's graph on the host in f32. Not a general ONNX runtime: it covers exactly the ops a // Whisper encoder uses (measured on both exports -- 447 nodes / 11 op types for whisper-small, 1031 / 10 for large-v3), and // it refuses an op it does not implement *by name* rather than producing a plausible activation from a silently skipped node. diff --git a/npu/audio_parallel.h b/npu/audio_parallel.h new file mode 100644 index 000000000..30ebfc05c --- /dev/null +++ b/npu/audio_parallel.h @@ -0,0 +1,81 @@ +// 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. +// +// audio_parallel.h -- the fork-join the host audio executor splits independent work over (engine#1931). +// +// Deterministic by construction: unit u always runs fn(u) exactly once, and a unit writes only its own slice of the +// output, so every element's arithmetic is unchanged by the thread count and the result is bit-identical to the serial +// loop. That is the property a reference implementation has to keep, and it is asserted, not assumed: +// tests/audio_parallel_test.cpp compares the executor's MatMul byte for byte across thread counts. +// +// A unit that throws does not take the process down. A std::thread whose function lets an exception escape calls +// std::terminate, so each worker catches, every worker is joined, and the first exception (lowest worker index, which +// makes the choice repeatable) is rethrown on the calling thread, where the executor turns it into a refusal with a +// reason like every other failure. +// +// Header-only and free of the executor's types, so the test needs no pack and no device. + +#pragma once + +#include +#include +#include +#include +#include + +namespace onebit::npu::pack_import { + +// How many threads the audio executor may use. ONEBIT_AUDIO_THREADS=N sets it; anything that is not a positive number +// means 1 (serial, the control for measuring the split), so a typo cannot turn into "as many as the parse produced". +// Unset: the hardware's count, at most 8. +inline size_t audio_thread_count() { + if (const char* env = std::getenv("ONEBIT_AUDIO_THREADS")) { + const long n = std::strtol(env, nullptr, 10); + return n > 0 ? size_t(n) : 1; + } + const size_t hw = std::thread::hardware_concurrency(); + return hw == 0 ? 1 : (hw > 8 ? 8 : hw); +} + +// Run fn(0) .. fn(units-1), each exactly once, on at most `want` threads. Worker t takes units t, t+want, t+2*want, ... +template +void parallel_units(size_t units, size_t want, Fn&& fn) { + if (want > units) want = units; + if (want <= 1) { + for (size_t u = 0; u < units; ++u) fn(u); + return; + } + std::vector failed(want); + std::vector pool; + pool.reserve(want); + for (size_t t = 0; t < want; ++t) + pool.emplace_back([&, t]() { + try { + for (size_t u = t; u < units; u += want) fn(u); + } catch (...) { + failed[t] = std::current_exception(); + } + }); + for (auto& th : pool) th.join(); + for (const auto& ep : failed) + if (ep) std::rethrow_exception(ep); +} + +template +void parallel_units(size_t units, Fn&& fn) { + parallel_units(units, audio_thread_count(), static_cast(fn)); +} + +} // namespace onebit::npu::pack_import diff --git a/npu/forward/include/audio_mel.h b/npu/forward/include/audio_mel.h index 2f7919891..bb0b2d8f8 100644 --- a/npu/forward/include/audio_mel.h +++ b/npu/forward/include/audio_mel.h @@ -1,6 +1,18 @@ // 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. +// // audio_mel.h -- the log-mel front-end: hann-windowed STFT (n_fft 400, hop 160, reflect padding), power spectrum, a // slaney-normalised mel filterbank, then log10 and the (max - 8, +4)/4 scaling a Whisper encoder's input expects. // diff --git a/npu/forward/include/audio_ops.h b/npu/forward/include/audio_ops.h index f8cde2e74..8fd3dfef6 100644 --- a/npu/forward/include/audio_ops.h +++ b/npu/forward/include/audio_ops.h @@ -1,6 +1,18 @@ // 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. +// // audio_ops.h -- the four f32 host kernels a Whisper encoder front-end is made of: a temporal conv with stride and zero // padding, layer norm, erf-GELU and a row softmax. Header-only, no dependencies. // diff --git a/npu/onnx_graph.h b/npu/onnx_graph.h index db3a8157a..87acc99cd 100644 --- a/npu/onnx_graph.h +++ b/npu/onnx_graph.h @@ -1,6 +1,18 @@ // 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. +// // onnx_graph.h -- a *reader* of an ONNX graph, as opposed to the census walk in npu/pack_import.cpp, which keeps only // the MatMul-family nodes because that is all a quantised LLM pack needs. An audio pack (the Whisper encoder) is a plain // f32 graph: 32 transformer blocks of Conv/MatMul/Add/LayerNormalization/Gelu/Softmax/Transpose/Mul, and to execute it the diff --git a/tests/audio_exec_test.cpp b/tests/audio_exec_test.cpp index 1a63a6ee0..0b923244a 100644 --- a/tests/audio_exec_test.cpp +++ b/tests/audio_exec_test.cpp @@ -1,6 +1,18 @@ // 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. +// // audio_exec_test.cpp -- run an audio pack's encoder through the engine's own host executor and compare it against // onnxruntime, on an input both sides can build from a formula (no binary fixture in the repo). // diff --git a/tests/audio_mel_test.cpp b/tests/audio_mel_test.cpp index 70d7e823b..e7c368443 100644 --- a/tests/audio_mel_test.cpp +++ b/tests/audio_mel_test.cpp @@ -1,6 +1,18 @@ // 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. +// // audio_mel_test.cpp -- the log-mel front-end against the reference extractor's numbers. // // Input: 30 s at 16 kHz of three summed sinusoids (440, 1250, 3100 Hz), built from a formula so both sides use the same diff --git a/tests/audio_ops_test.cpp b/tests/audio_ops_test.cpp index 68f5a6087..9107a8c53 100644 --- a/tests/audio_ops_test.cpp +++ b/tests/audio_ops_test.cpp @@ -1,6 +1,18 @@ // 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. +// // audio_ops_test.cpp -- known-answer checks for npu/forward/include/audio_ops.h. No container, no device, no pack: every // expected value was computed in double in Python from the same formulas the Whisper encoder reference uses // (tools/whisper_encoder_ref.py), so this binary runs wherever CI builds, which is the point. diff --git a/tests/audio_parallel_test.cpp b/tests/audio_parallel_test.cpp new file mode 100644 index 000000000..93b71a551 --- /dev/null +++ b/tests/audio_parallel_test.cpp @@ -0,0 +1,198 @@ +// 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. +// +// audio_parallel_test.cpp -- the two claims npu/audio_parallel.h makes, as assertions (engine#1931 review): +// +// 1. the executor's MatMul is BIT-identical whatever ONEBIT_AUDIO_THREADS says. A tolerance cannot see this, and the +// pack-gated comparison against onnxruntime runs at one thread count, so here the same MatMul runs through +// run_audio_graph() at 1, 2, 3, 7, 8 and 64 threads and unset, and the outputs are compared with memcmp; +// 2. a unit that throws is an exception on the calling thread after every worker has been joined, not std::terminate. +// +// The graph is built in memory (one MatMul against an inline initializer), so this needs no pack, no file and no device. + +#include +#include +#include +#include +#include +#include +#include +#include + +#include "audio_exec.h" +#include "audio_parallel.h" +#include "onnx_graph.h" + +namespace { + +int g_failures = 0; + +void check(bool cond, const std::string& what) { + if (!cond) { + std::printf("FAIL %s\n", what.c_str()); + ++g_failures; + } +} + +using namespace onebit::npu::pack_import; + +// y = x . W for x (1, M, K) and W (K, N): M rows, which is what the executor splits over. +std::vector matmul_at(const char* threads, int64_t M, int64_t K, int64_t N, const std::vector& x, + const std::vector& w) { + if (threads != nullptr) setenv("ONEBIT_AUDIO_THREADS", threads, 1); + else unsetenv("ONEBIT_AUDIO_THREADS"); + OnnxGraphInfo g; + g.dir = "."; + g.inputs = {"x"}; + g.outputs = {"y"}; + OnnxInit wi; + wi.name = "W"; + wi.dtype = 1; + wi.dims = {K, N}; + wi.raw = reinterpret_cast(w.data()); + wi.raw_len = w.size() * sizeof(float); + g.initializers = {"W"}; + g.init_tensors.push_back(wi); + OnnxNodeInfo n; + n.op = "MatMul"; + n.name = "mm"; + n.ins = {"x", "W"}; + n.outs = {"y"}; + g.nodes.push_back(n); + AudioRunResult out; + std::string why; + if (!run_audio_graph(g, x.data(), M, K, "y", &out, &why)) { + std::printf("FAIL MatMul at ONEBIT_AUDIO_THREADS=%s refused: %s\n", threads ? threads : "(unset)", why.c_str()); + ++g_failures; + return {}; + } + check(out.dims == std::vector({1, M, N}), "MatMul output is (1, M, N)"); + return out.data; +} + +void test_bit_identity() { + // 37 rows: divisible by none of the thread counts below, so the last worker's share is short at every one of them. + const int64_t M = 37, K = 53, N = 29; + std::vector x(size_t(M * K)), w(size_t(K * N)); + for (size_t i = 0; i < x.size(); ++i) x[i] = float(std::sin(0.37 * double(i + 1)) * 3.0); + for (size_t i = 0; i < w.size(); ++i) w[i] = float(std::cos(0.11 * double(i + 3)) * 0.7 + 1e-3 * double(i % 7)); + for (size_t i = 0; i < x.size(); i += 5) x[i] = 0.0f; // the zero shortcut is on the path + for (int64_t k = 0; k < K; ++k) x[size_t(11 * K + k)] = 0.0f; // and one whole row of zeros + + const std::vector serial = matmul_at("1", M, K, N, x, w); + check(serial.size() == size_t(M * N), "the serial MatMul produced M*N values"); + if (serial.size() != size_t(M * N)) return; + + // The serial result is the arithmetic the split must not change: an ordered sum over k, per element, in float. + bool same_as_plain_loop = true; + for (int64_t i = 0; i < M && same_as_plain_loop; ++i) + for (int64_t j = 0; j < N; ++j) { + float acc = 0.0f; + for (int64_t k = 0; k < K; ++k) { + const float av = x[size_t(i * K + k)]; + if (av == 0.0f) continue; + acc += av * w[size_t(k * N + j)]; + } + if (std::memcmp(&acc, &serial[size_t(i * N + j)], sizeof acc) != 0) { + same_as_plain_loop = false; + break; + } + } + check(same_as_plain_loop, "the serial MatMul is the ordered per-element sum, bit for bit"); + bool zero_row = true; + for (int64_t j = 0; j < N; ++j) zero_row = zero_row && serial[size_t(11 * N + j)] == 0.0f; + check(zero_row, "a row of zeros gives a row of zeros"); + + for (const char* threads : {"2", "3", "7", "8", "64", static_cast(nullptr)}) { + const std::vector got = matmul_at(threads, M, K, N, x, w); + const std::string name = threads ? threads : "(unset)"; + check(got.size() == serial.size() && std::memcmp(got.data(), serial.data(), serial.size() * sizeof(float)) == 0, + "MatMul at ONEBIT_AUDIO_THREADS=" + name + " is bit-identical to the serial result"); + } + // What is not a thread count means serial, and the answer is still the same one. + for (const char* threads : {"0", "-4", "many", ""}) { + setenv("ONEBIT_AUDIO_THREADS", threads, 1); + check(audio_thread_count() == 1, std::string("ONEBIT_AUDIO_THREADS='") + threads + "' means one thread"); + const std::vector got = matmul_at(threads, M, K, N, x, w); + check(got.size() == serial.size() && std::memcmp(got.data(), serial.data(), serial.size() * sizeof(float)) == 0, + std::string("and MatMul at ONEBIT_AUDIO_THREADS='") + threads + "' is the serial result"); + } + unsetenv("ONEBIT_AUDIO_THREADS"); +} + +void test_every_unit_once() { + for (size_t units : {size_t(0), size_t(1), size_t(2), size_t(7), size_t(37), size_t(1000)}) + for (size_t want : {size_t(0), size_t(1), size_t(2), size_t(3), size_t(8), size_t(64), size_t(5000)}) { + std::vector> seen(units); + for (auto& s : seen) s = 0; + parallel_units(units, want, [&](size_t u) { seen[u].fetch_add(1); }); + bool once = true; + for (auto& s : seen) once = once && s.load() == 1; + check(once, "every one of " + std::to_string(units) + " units runs exactly once on " + std::to_string(want) + + " thread(s)"); + } +} + +void test_a_throwing_unit() { + // If the exception escaped a worker, the process would call std::terminate here and the test would die by signal. + for (size_t want : {size_t(1), size_t(2), size_t(4), size_t(8)}) { + std::atomic ran{0}; + bool caught = false; + std::string what; + try { + parallel_units(40, want, [&](size_t u) { + ran.fetch_add(1); + if (u == 5) throw std::runtime_error("unit 5 refused"); + }); + } catch (const std::runtime_error& e) { + caught = true; + what = e.what(); + } + check(caught && what == "unit 5 refused", + "a unit that throws on " + std::to_string(want) + " thread(s) is an exception the caller catches"); + // the other workers ran to the end and were joined before the rethrow: only the throwing worker stopped early + if (want > 1) { + const int stopped_worker_left = int((40 - 1 - 5) / want); // units after 5 in unit 5's worker + check(ran.load() == 40 - stopped_worker_left, + "the other workers finish and are joined before the caller sees it (" + std::to_string(want) + " threads)"); + } + } + // Two workers throw: the caller gets the lower worker's, every time. + for (int rep = 0; rep < 20; ++rep) { + std::string what; + try { + parallel_units(8, 4, [&](size_t u) { + if (u == 1) throw std::runtime_error("worker 1"); + if (u == 3) throw std::runtime_error("worker 3"); + }); + } catch (const std::runtime_error& e) { + what = e.what(); + } + if (what != "worker 1") { + check(false, "when two workers throw, the lower worker's exception is the one rethrown (got '" + what + "')"); + break; + } + } +} + +} // namespace + +int main() { + test_bit_identity(); + test_every_unit_once(); + test_a_throwing_unit(); + if (g_failures == 0) std::printf("audio_parallel: ok\n"); + return g_failures == 0 ? 0 : 1; +} diff --git a/tests/onnx_graph_test.cpp b/tests/onnx_graph_test.cpp index c2a2724a4..2d6234dc5 100644 --- a/tests/onnx_graph_test.cpp +++ b/tests/onnx_graph_test.cpp @@ -1,6 +1,18 @@ // 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. +// // onnx_graph_test.cpp -- the audio pack's graph, read through the engine and checked for the property the executor will // depend on: every node's inputs are graph inputs, initializers, or outputs of an *earlier* node. Skips (77) without // ONEBIT_TEST_WHISPER naming the encoder's .onnx. diff --git a/tools/whisper_transcribe.cpp b/tools/whisper_transcribe.cpp index 004c74d42..119a17131 100644 --- a/tools/whisper_transcribe.cpp +++ b/tools/whisper_transcribe.cpp @@ -1,6 +1,18 @@ // 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. +// // whisper_transcribe.cpp -- the whole host pipeline: wav -> log-mel -> encoder -> greedy decode -> token ids. // // ./whisper_transcribe [expected-ids-csv ...] @@ -96,6 +108,10 @@ int main(int argc, char** argv) { using namespace onebit::npu::pack_import; std::string why; + // The spans the last line reports, each named for what it contains and nothing else: wav (reading the file), mel + // (the log-mel front-end), load (parsing both graphs), encoder, decode. `wall` is their sum, from before the file + // is opened. + const auto t_begin = std::chrono::steady_clock::now(); std::vector audio = read_wav_16k_mono(argv[3], &why); if (audio.empty()) { std::fprintf(stderr, "wav: %s\n", why.c_str()); @@ -121,6 +137,7 @@ int main(int argc, char** argv) { std::fprintf(stderr, "graph: %s\n", e.what()); return 1; } + const auto t_load = std::chrono::steady_clock::now(); AudioRunResult xa; if (!run_audio_graph(enc, mel.data(), cfg.bins, cfg.frames, "", &xa, &why)) { std::fprintf(stderr, "encoder: %s\n", why.c_str()); @@ -158,9 +175,10 @@ int main(int argc, char** argv) { const auto secs = [](std::chrono::steady_clock::time_point a, std::chrono::steady_clock::time_point b) { return std::chrono::duration(b - a).count(); }; - std::printf("time: mel+load %.3f s, encoder %.3f s, decode %.3f s for %zu token(s) -> %.3f s/token, wall %.3f s\n", - secs(t_start, t_mel), secs(t_mel, t_enc), secs(t_enc, t_dec), ids.size(), - ids.empty() ? 0.0 : secs(t_enc, t_dec) / double(ids.size()), secs(t_start, t_dec)); + std::printf("time: wav %.3f s, mel %.3f s, load %.3f s, encoder %.3f s, decode %.3f s for %zu token(s) -> %.3f s/token, " + "wall %.3f s\n", + secs(t_begin, t_start), secs(t_start, t_mel), secs(t_mel, t_load), secs(t_load, t_enc), secs(t_enc, t_dec), + ids.size(), ids.empty() ? 0.0 : secs(t_enc, t_dec) / double(ids.size()), secs(t_begin, t_dec)); std::printf("ids:"); for (int64_t v : ids) std::printf(" %lld", (long long)v); std::printf("\n");