laya: an encoder registry; serve --laya runs the encoder on the NPU through the private add-on (ONEBIT_LAYA_DEVICE=npu)

.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

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)

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<std::string> 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<LayaDaemon>(std::vector<std::string>{
@@ -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=<npu-kernels checkout>, docs/laya.md)");
             const std::string ckpt = laya_checkpoint(o.laya_model);
             scorer = std::make_unique<onebit::laya::Scorer>();
             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 {

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=<npu-kernels checkout> ...
+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.

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=<npu-kernels checkout>". 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.
 

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 <map>
+
+namespace onebit::laya {
+
+namespace {
+std::map<std::string, EncoderFactory>& registry() {
+    static std::map<std::string, EncoderFactory> 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<std::string> encoder_names() {
+    std::vector<std::string> names;
+    for (const auto& [name, f] : registry()) names.push_back(name);
+    return names;
+}
+
+}  // namespace onebit::laya

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 <functional>
+#include <string>
+#include <vector>
+
+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<bool(float* h, int rows, int n)>;
+
+// model_dir (the checkpoint Scorer::load read) -> an encoder; on failure an empty function and
+// the reason in *err.
+using EncoderFactory = std::function<EncoderFn(const std::string& model_dir, std::string* err)>;
+
+// 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<std::string> encoder_names();
+
+}  // namespace onebit::laya

laya/scorer.cpp

@@ -399,7 +399,21 @@ bool Scorer::score(const std::string& state, const std::vector<Question>& questi
 
     std::vector<float> hn((size_t)N * L * D), qkv((size_t)N * L * 3072), attn_out((size_t)N * L * D);
     std::vector<float> 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<float> 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;

laya/scorer.h

@@ -25,6 +25,7 @@
 #pragma once
 
 #include <cstdint>
+#include <functional>
 #include <memory>
 #include <string>
 #include <utility>
@@ -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<bool(float* h, int rows, int n)> 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<float> temperature_;                        // [3]
     std::vector<std::pair<std::string, float>> temperature_by_options_;
 
+    std::function<bool(float* h, int rows, int n)> encoder_;
+    long encoded_ = 0, declined_ = 0;
     std::unique_ptr<onebit::npu::Tokenizer> tok_;  // loaded once: reading tokenizer.json costs ~50 ms
     std::string err_;
 };

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 <cstdio>
+#include <string>
+
+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<std::string>{"npu"});
+    return failures ? 1 : 0;
+}

tests/laya_gguf_route.sh

@@ -25,7 +25,8 @@
 #   - each class reaches its device, with X-1bit-Route "<class> <confidence> <device>",
 #   - 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" \