diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index fd31d9e75..09a45fb31 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -49,6 +49,6 @@ jobs: - name: Configure run: cmake -B build -G Ninja -DCMAKE_BUILD_TYPE=Release -DCMAKE_C_COMPILER=gcc-15 -DCMAKE_CXX_COMPILER=g++-15 - name: Build - run: cmake --build build --target onebit npu_full_elf_test npu_pack_test npu_q4nx_test npu_private_route_test tokenizer_test think_split_test gguf_meta_test -j + run: cmake --build build --target onebit npu_full_elf_test npu_pack_test npu_q4nx_test npu_private_route_test tokenizer_test think_split_test gguf_meta_test laya_encoder_test -j - name: Test run: ctest --test-dir build --output-on-failure diff --git a/CMakeLists.txt b/CMakeLists.txt index 3ee7aef73..36ba134e2 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -297,7 +297,7 @@ file(WRITE ${CMAKE_BINARY_DIR}/generated/route_policy_default.h "// generated from config/route-policy.json by CMakeLists.txt\n#pragma once\n" "static const char kDefaultRoutePolicyJson[] = R\"ONEBITPOLICY(${ONEBIT_ROUTE_POLICY_JSON})ONEBITPOLICY\";\n") set_property(DIRECTORY APPEND PROPERTY CMAKE_CONFIGURE_DEPENDS ${CMAKE_SOURCE_DIR}/config/route-policy.json) -add_library(onebit_laya STATIC laya/scorer.cpp laya/safetensors.cpp laya/route.cpp laya/policy.cpp) +add_library(onebit_laya STATIC laya/scorer.cpp laya/safetensors.cpp laya/route.cpp laya/policy.cpp laya/encoder.cpp) target_include_directories(onebit_laya PUBLIC laya PRIVATE ${CMAKE_BINARY_DIR}/generated) target_link_libraries(onebit_laya PUBLIC onebit_npu nlohmann_json::nlohmann_json) # The engine binary gets `1bit route` and per-request routing through the scorer. @@ -313,6 +313,9 @@ if(ONEBIT_LAYA_GGML) endif() add_executable(laya_gate tests/laya_gate.cpp) +add_executable(laya_encoder_test tests/laya_encoder_test.cpp) +target_link_libraries(laya_encoder_test PRIVATE onebit_laya) +add_test(NAME laya_encoder COMMAND laya_encoder_test) add_executable(laya_classify_eval tests/laya_classify_eval.cpp) target_link_libraries(laya_classify_eval PRIVATE onebit_laya nlohmann_json::nlohmann_json) target_link_libraries(laya_gate PRIVATE onebit_laya nlohmann_json::nlohmann_json) diff --git a/app/serve.cpp b/app/serve.cpp index 54c149a0d..428462c2f 100644 --- a/app/serve.cpp +++ b/app/serve.cpp @@ -79,6 +79,7 @@ #include "unified.h" #endif #ifdef ONEBIT_LAYA +#include "encoder.h" #include "policy.h" #include "route.h" #include "scorer.h" @@ -1081,12 +1082,14 @@ int serve_child(const Options& given) { if (o.adaptive || !o.role.empty() || !o.prefill_device.empty() || o.lean) throw std::runtime_error("--laya does not combine with --adaptive, --embedding/--reranking, --prefill-device or --lean"); policy = onebit::laya::load_route_policy(o.route_policy); - const std::string gguf = laya_gguf(o.laya_model), bin = laya_ggml_bin(); + // ONEBIT_LAYA_DEVICE=npu: the C++ scorer with its encoder on the NPU, from a private add-on + // (laya/encoder.h, docs/laya.md); else the GGUF scorer, on HRX0 unless it names another + const char* dev = std::getenv("ONEBIT_LAYA_DEVICE"); + const bool laya_npu = dev && std::string(dev) == "npu"; + const std::string gguf = laya_npu ? "" : laya_gguf(o.laya_model), bin = laya_npu ? "" : laya_ggml_bin(); #if defined(__linux__) if (!gguf.empty() && !bin.empty()) { - // the GGUF scorer on HRX0 (ONEBIT_LAYA_DEVICE picks another: cpu), with the HSA runtime HRX - // needs (docs/hrx.md) - const char* dev = std::getenv("ONEBIT_LAYA_DEVICE"); + // the GGUF scorer, with the HSA runtime HRX needs (docs/hrx.md) std::vector env; if (const std::string hsa = hrx_libhsa(o.hrx_libhsa); !hsa.empty()) env.push_back("IREE_HAL_AMDGPU_LIBHSA_PATH=" + hsa); laya_daemon = std::make_unique(std::vector{ @@ -1101,11 +1104,21 @@ int serve_child(const Options& given) { if (!laya_daemon) #endif { + const onebit::laya::EncoderFactory* f = laya_npu ? onebit::laya::find_encoder("npu") : nullptr; + if (laya_npu && !f) + throw std::runtime_error("laya: ONEBIT_LAYA_DEVICE=npu needs a build with the NPU add-on " + "(-DONEBIT_NPU_PRIVATE=, docs/laya.md)"); const std::string ckpt = laya_checkpoint(o.laya_model); scorer = std::make_unique(); if (!scorer->load(ckpt)) throw std::runtime_error("laya: " + scorer->error()); - std::fprintf(stderr, "1bit serve: laya: checkpoint %s, policy %s\n", ckpt.c_str(), - o.route_policy.empty() ? "built in" : o.route_policy.c_str()); + if (f) { + std::string err; + onebit::laya::EncoderFn fn = (*f)(ckpt, &err); + if (!fn) throw std::runtime_error("laya: the NPU encoder: " + err); + scorer->set_encoder(std::move(fn)); + } + std::fprintf(stderr, "1bit serve: laya: checkpoint %s%s, policy %s\n", ckpt.c_str(), + laya_npu ? " (encoder on the NPU)" : "", o.route_policy.empty() ? "built in" : o.route_policy.c_str()); } for (const std::string& dev : kGgufDevices) { try { diff --git a/docs/laya.md b/docs/laya.md index 2710258cd..91115066d 100644 --- a/docs/laya.md +++ b/docs/laya.md @@ -222,3 +222,26 @@ typed-decisions for all of them; Hindi is the weak spot either way. ggmlc declares the MIT licence (`pyproject.toml`, README) but its repository carries no LICENSE file yet; the GGUFs are Apache-2.0 like the checkpoints. + +## The encoder on the NPU (private add-on) + +`laya/encoder.h` lets an add-on run the C++ scorer's 28 encoder layers somewhere else: it +registers an encoder by name, and `Scorer::set_encoder` hands it each sequence's embedded rows +(the tokenizer, embeddings and classifier head stay on the CPU). The engine registers none. The +private NPU add-on (docs/npu.md, "Private routes"; its kernels live in the closed-source +`npu-kernels` repository) registers `npu`, and + +```sh +cmake -B build -DONEBIT_NPU=ON -DONEBIT_NPU_PRIVATE= ... +ONEBIT_LAYA_DEVICE=npu 1bit serve -m MODEL.gguf --laya +``` + +runs the encoder on the NPU instead of the GGUF daemon. It shares the NPU with other users +through the engine's device lock: while another process holds it (an NPU `1bit serve`), the +encoder declines and that decision's layers run on the CPU. Without the add-on, +`ONEBIT_LAYA_DEVICE=npu` exits saying how to build it (checked by `tests/laya_gguf_route.sh`). + +**Measured** (Strix Halo, `tests/laya_route_cases.json`): 194/200 = 97.0% class accuracy (the +CPU scorer: 95.5%), about 105 ms a warm decision through `serve` (encoder ~50 ms, the CPU head +~41 ms), about 500 ms when the NPU is busy and the CPU runs the layers. Slower than the GGUF +scorer on HRX today; it keeps the GPU free for the model being served. diff --git a/docs/npu.md b/docs/npu.md index 9401cedb5..a79a5a30f 100644 --- a/docs/npu.md +++ b/docs/npu.md @@ -342,6 +342,9 @@ The engine keeps only the hook (`npu/private_route.h`): Qwen3.6-35B-A3B NPU route is not part of this build; build with -DONEBIT_NPU_PRIVATE=". The fast lane is unaffected. +The add-on may also register an encoder for the Laya router (`laya/encoder.h`, docs/laya.md: +`ONEBIT_LAYA_DEVICE=npu`). + `tests/npu_private_route_test.cpp` (ctest `npu_private_route`, in CI) checks the registry and that message; the add-on brings its own tests. diff --git a/laya/encoder.cpp b/laya/encoder.cpp new file mode 100644 index 000000000..a47f27227 --- /dev/null +++ b/laya/encoder.cpp @@ -0,0 +1,42 @@ +// 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. + +#include "encoder.h" + +#include + +namespace onebit::laya { + +namespace { +std::map& registry() { + static std::map r; + return r; +} +} // namespace + +void register_encoder(const std::string& name, EncoderFactory factory) { registry()[name] = std::move(factory); } + +const EncoderFactory* find_encoder(const std::string& name) { + const auto it = registry().find(name); + return it == registry().end() ? nullptr : &it->second; +} + +std::vector encoder_names() { + std::vector names; + for (const auto& [name, f] : registry()) names.push_back(name); + return names; +} + +} // namespace onebit::laya diff --git a/laya/encoder.h b/laya/encoder.h new file mode 100644 index 000000000..b053b2636 --- /dev/null +++ b/laya/encoder.h @@ -0,0 +1,44 @@ +// 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. + +// laya/encoder.h — another place for the scorer's encoder (its 28 ModernBERT layers): the scorer +// hands each sequence's embedded rows to a registered encoder and runs the layers on the CPU only +// when the encoder declines. The engine registers none; a private NPU add-on (docs/npu.md, +// "Private routes") registers "npu" from register_private_addon(), and `1bit serve --laya` with +// ONEBIT_LAYA_DEVICE=npu uses it (docs/laya.md). +#pragma once + +#include +#include +#include + +namespace onebit::laya { + +// Runs the encoder layers on h [rows][1024] in place: n real tokens, the rows after them padding +// (masked as keys). Leaves h as the last layer's output, before the final norm. false: declined +// (h unchanged; the scorer runs the layers itself). +using EncoderFn = std::function; + +// model_dir (the checkpoint Scorer::load read) -> an encoder; on failure an empty function and +// the reason in *err. +using EncoderFactory = std::function; + +// A later registration under the same name replaces the earlier one. +void register_encoder(const std::string& name, EncoderFactory factory); +// nullptr when nothing is registered under name. +const EncoderFactory* find_encoder(const std::string& name); +std::vector encoder_names(); + +} // namespace onebit::laya diff --git a/laya/scorer.cpp b/laya/scorer.cpp index a4220edbe..e46f9e2e2 100644 --- a/laya/scorer.cpp +++ b/laya/scorer.cpp @@ -399,7 +399,21 @@ bool Scorer::score(const std::string& state, const std::vector& questi std::vector hn((size_t)N * L * D), qkv((size_t)N * L * 3072), attn_out((size_t)N * L * D); std::vector q((size_t)N * NHEAD * L * HD), k((size_t)N * NHEAD * L * HD), v((size_t)N * NHEAD * L * HD); - for (int li = 0; li < NLAYER; ++li) { + // a set encoder (laya/encoder.h) takes each sequence's rows; if it declines one, the CPU runs all + bool encoded = false; + if (encoder_) { + std::vector out(h); + encoded = true; + for (int b = 0; b < N && encoded; ++b) + encoded = encoder_(out.data() + (size_t)b * L * D, L, (int)seqs[b].size()); + if (encoded) { + h.swap(out); + encoded_ += N; + } else { + declined_ += N; + } + } + for (int li = 0; li < NLAYER && !encoded; ++li) { LayerW& Lw = layers_[li]; bool is_full = (li % 3 == 0); const auto& cs = is_full ? full_cs : slid_cs; diff --git a/laya/scorer.h b/laya/scorer.h index afdb06410..6c6bce823 100644 --- a/laya/scorer.h +++ b/laya/scorer.h @@ -25,6 +25,7 @@ #pragma once #include +#include #include #include #include @@ -80,6 +81,13 @@ class Scorer { const std::string& error() const { return err_; } + // Runs the encoder layers through fn (laya/encoder.h) from now on; an empty fn goes back to the + // CPU. A call where fn declines any sequence runs the whole batch on the CPU. + void set_encoder(std::function fn) { encoder_ = std::move(fn); } + // Sequences the set encoder ran, and sequences it declined (run on the CPU instead). + long encoded() const { return encoded_; } + long declined() const { return declined_; } + private: // config int max_len_ = 512; @@ -110,6 +118,8 @@ class Scorer { std::vector temperature_; // [3] std::vector> temperature_by_options_; + std::function encoder_; + long encoded_ = 0, declined_ = 0; std::unique_ptr tok_; // loaded once: reading tokenizer.json costs ~50 ms std::string err_; }; diff --git a/tests/laya_encoder_test.cpp b/tests/laya_encoder_test.cpp new file mode 100644 index 000000000..32cd87894 --- /dev/null +++ b/tests/laya_encoder_test.cpp @@ -0,0 +1,56 @@ +// 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 encoder registry (laya/encoder.h): what a private add-on registers is found by name, +// replaced by a later registration, and hands back the encoder its factory builds. +#include "encoder.h" + +#include +#include + +namespace { +int failures = 0; +void check(const char* what, bool ok) { + std::printf("%s %s\n", ok ? "PASS" : "FAIL", what); + if (!ok) failures++; +} +} // namespace + +int main() { + using namespace onebit::laya; + check("nothing registered: npu not found", find_encoder("npu") == nullptr); + register_encoder("npu", [](const std::string& dir, std::string* err) -> EncoderFn { + if (dir.empty()) { + *err = "no checkpoint"; + return {}; + } + return [](float* h, int rows, int n) { + for (int i = 0; i < rows * 1024; i++) h[i] += 1.f; + return n <= rows; + }; + }); + const EncoderFactory* f = find_encoder("npu"); + check("registered: npu found", f != nullptr); + std::string err; + check("factory failure: empty encoder and a reason", f && !(*f)("", &err) && err == "no checkpoint"); + EncoderFn fn = f ? (*f)("ckpt", &err) : EncoderFn{}; + float h[2 * 1024] = {}; + check("the encoder runs", fn && fn(h, 2, 1) && h[0] == 1.f && h[2047] == 1.f); + check("it may decline", fn && !fn(h, 1, 2)); + register_encoder("npu", [](const std::string&, std::string*) -> EncoderFn { return {}; }); + check("a later registration replaces it", !(*find_encoder("npu"))("ckpt", &err)); + check("names", encoder_names() == std::vector{"npu"}); + return failures ? 1 : 0; +} diff --git a/tests/laya_gguf_route.sh b/tests/laya_gguf_route.sh index b4e58007c..7f16bb0c1 100755 --- a/tests/laya_gguf_route.sh +++ b/tests/laya_gguf_route.sh @@ -25,7 +25,8 @@ # - each class reaches its device, with X-1bit-Route " ", # - the daemon gets serve's question: the three non-long_doc classes, the request as state, # - a request of 1024 characters or more is long_doc without asking the daemon, -# - --laya-model FILE.gguf without ggmlc's laya fails and says how to build it. +# - --laya-model FILE.gguf without ggmlc's laya fails and says how to build it, +# - ONEBIT_LAYA_DEVICE=npu without the NPU add-on (laya/encoder.h) fails and says how to build it. # # usage: tests/laya_gguf_route.sh path/to/1bit set -uo pipefail @@ -66,6 +67,9 @@ nobin=$(env -u ONEBIT_LAYA_GGML "$bin" serve -m "$scratch/tiny.gguf" --device au --laya-model "$scratch/laya/gguf/laya_typed_decisions_q8_0.gguf" --port "$port" \ --llama-server "$here/fake_backend_route.py" 2>&1) check "--laya-model FILE.gguf without ggmlc's laya says how to build it" '[[ "$nobin" == *ONEBIT_LAYA_GGML* ]]' +nonpu=$(ONEBIT_LAYA_DEVICE=npu ONEBIT_LAYA_MODEL="$scratch/laya" "$bin" serve -m "$scratch/tiny.gguf" --device auto --laya \ + --port "$port" --llama-server "$here/fake_backend_route.py" 2>&1) +check "ONEBIT_LAYA_DEVICE=npu without the NPU add-on says how to build it" '[[ "$nonpu" == *ONEBIT_NPU_PRIVATE* ]]' LAYA_LOG="$scratch/daemon.log" ONEBIT_LAYA_GGML="$scratch/laya-daemon.py" ONEBIT_LAYA_MODEL="$scratch/laya" \ "$bin" serve -m "$scratch/tiny.gguf" --device auto --laya --route-policy "$here/route-policy-e2e.json" \