diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index cc062e68..7743f62c 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -29,6 +29,6 @@ jobs: - name: Configure run: cmake -B build -G Ninja -DCMAKE_BUILD_TYPE=Release - name: Build - run: cmake --build build --target onebit npu_full_elf_test npu_pack_test npu_q4nx_test tokenizer_test -j + run: cmake --build build --target onebit npu_full_elf_test npu_pack_test npu_q4nx_test npu_lax_test tokenizer_test -j - name: Test run: ctest --test-dir build --output-on-failure diff --git a/.gitmodules b/.gitmodules index 4a305a2d..66e6c39c 100644 --- a/.gitmodules +++ b/.gitmodules @@ -31,6 +31,12 @@ branch = main # ZINC (docs/zinc.md): upstream zolotukhin/zinc, built by scripts/build-zinc.sh; # .github/workflows/bump-zinc.yml keeps it current. +# The 35B MoE whole-layer decode (docs/npu-lax.md): the open kernels, MIT, built by +# scripts/build-lax.sh; the branch carries our fixes and optimisations of the lax design. +[submodule "third_party/OpenFlowLM-Next"] + path = third_party/OpenFlowLM-Next + url = https://github.com/bong-water-water-bong/OpenFlowLM-Next.git + branch = 1bit/lax-35b [submodule "third_party/zinc"] path = third_party/zinc url = https://github.com/zolotukhin/zinc.git diff --git a/CMakeLists.txt b/CMakeLists.txt index a68bf29f..61dd87ea 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -182,7 +182,11 @@ option(ONEBIT_NPU_HOST "Build the NPU host library and its tests" ${_onebit_npu_ if(ONEBIT_NPU_HOST) find_package(PkgConfig REQUIRED) pkg_check_modules(PCRE2 REQUIRED IMPORTED_TARGET libpcre2-8) -add_library(onebit_npu STATIC npu/full_elf.cpp npu/model.cpp npu/pack.cpp npu/q4nx.cpp npu/tokenizer.cpp) +add_library(onebit_npu STATIC npu/full_elf.cpp npu/lax_elf.cpp npu/lax_pack.cpp npu/lax_stream.cpp npu/lax_turns.cpp npu/model.cpp npu/pack.cpp npu/q4nx.cpp + npu/tokenizer.cpp) +# The lax packer's q8 -> q4_1 re-quantization must round every float32 step on its own, +# as the reference packer's NumPy does (npu/lax_pack.h): no fused multiply-adds. +set_source_files_properties(npu/lax_pack.cpp PROPERTIES COMPILE_OPTIONS -ffp-contract=off) target_include_directories(onebit_npu PUBLIC npu) target_link_libraries(onebit_npu PUBLIC nlohmann_json::nlohmann_json PRIVATE PkgConfig::PCRE2) @@ -215,6 +219,16 @@ 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() +# The Qwen3.6-35B-A3B lax decode's host side (docs/npu-lax.md): CI checks the transforms +# against the reference packer's hashes; with the model, all 83 buffers against its table. +add_executable(npu_lax_test tests/npu_lax_test.cpp) +target_link_libraries(npu_lax_test PRIVATE onebit_npu) +add_test(NAME npu_lax_host COMMAND npu_lax_test) +set(ONEBIT_NPU_LAX_MODEL "" CACHE PATH "Qwen3.6-35B-A3B Q4NX directory, for the lax packer and decode tests") +if(ONEBIT_NPU_LAX_MODEL) + add_test(NAME npu_lax_pack_model + COMMAND npu_lax_test --model ${ONEBIT_NPU_LAX_MODEL} --sha256 ${CMAKE_SOURCE_DIR}/tests/golden/npu_lax/sha256.tsv) +endif() endif() # ONEBIT_NPU_HOST # The lane itself: needs XRT (the NPU runtime) and an XDNA NPU to run. @@ -227,10 +241,10 @@ if(ONEBIT_NPU) if(NOT EXISTS "${ONEBIT_XRT_ROOT}/include/xrt/xrt_device.h") message(FATAL_ERROR "ONEBIT_NPU needs XRT; none at ONEBIT_XRT_ROOT=${ONEBIT_XRT_ROOT}") endif() - add_library(onebit_npu_lane STATIC npu/lane.cpp npu/generate.cpp) + add_library(onebit_npu_lane STATIC npu/lane.cpp npu/generate.cpp npu/lax.cpp npu/lax_elf_kernels.cpp) target_include_directories(onebit_npu_lane SYSTEM PUBLIC ${ONEBIT_XRT_ROOT}/include) target_link_libraries(onebit_npu_lane PUBLIC onebit_npu ${ONEBIT_XRT_ROOT}/lib/libxrt_coreutil.so) - target_sources(onebit PRIVATE app/unified.cpp) + target_sources(onebit PRIVATE app/unified.cpp app/npu_lax.cpp) target_link_libraries(onebit PRIVATE onebit_npu_lane) target_compile_definitions(onebit PRIVATE ONEBIT_NPU) set_target_properties(onebit PROPERTIES BUILD_RPATH ${ONEBIT_XRT_ROOT}/lib INSTALL_RPATH ${ONEBIT_XRT_ROOT}/lib) @@ -241,6 +255,47 @@ if(ONEBIT_NPU) COMMAND ${CMAKE_SOURCE_DIR}/tests/npu_lane_e2e.sh $ ${ONEBIT_NPU_MODEL} ${ONEBIT_NPU_KERNELS} ${ONEBIT_NPU_REFERENCE}) endif() + # `1bit npu-lax`: parity on three positions against the fp64 reference, then one chat turn. + set(ONEBIT_NPU_LAX_KERNELS "" CACHE PATH "scripts/build-lax.sh's /kernels") + set(ONEBIT_NPU_LAX_REF "" CACHE PATH "make_decode.py --requant --tokens 3 output (xres.bin, ref_logits*.bin)") + if(ONEBIT_NPU_LAX_KERNELS) + # The full-attention stream's position patches, against the reference harness's table + # for the pinned build (tests/golden/npu_lax/lax_a_patches.tsv names its md5). + add_test(NAME npu_lax_patches + COMMAND npu_lax_test --insts ${ONEBIT_NPU_LAX_KERNELS}/lax_a/insts.bin + --patches ${CMAKE_SOURCE_DIR}/tests/golden/npu_lax/lax_a_patches.tsv) + # The full-ELF position sites of lax_a/insts.elf (tests/golden/npu_lax/lax_a_elf_sites.tsv + # names its md5); with the harness's full ELFs for the same build, every assembled ELF + # byte for byte. + set(ONEBIT_NPU_LAX_ELF_GOLDEN "" CACHE PATH "harness/full_elf.py lax output for the same kernel build (lax_init.elf, lxf.elf, axf_p.elf, ln.elf, lm.elf)") + set(_golden) + if(ONEBIT_NPU_LAX_ELF_GOLDEN) + set(_golden --golden ${ONEBIT_NPU_LAX_ELF_GOLDEN}) + endif() + add_test(NAME npu_lax_elf + COMMAND npu_lax_test --elf ${ONEBIT_NPU_LAX_KERNELS} + --sites ${CMAKE_SOURCE_DIR}/tests/golden/npu_lax/lax_a_elf_sites.tsv ${_golden}) + endif() + if(ONEBIT_NPU_LAX_MODEL AND ONEBIT_NPU_LAX_KERNELS AND ONEBIT_NPU_LAX_REF) + # Full ELFs (the default), and the classic xclbin path kept for A/B. + add_test(NAME npu_lax_e2e + COMMAND ${CMAKE_SOURCE_DIR}/tests/npu_lax_cpp.sh $ + ${ONEBIT_NPU_LAX_MODEL} ${ONEBIT_NPU_LAX_KERNELS} ${ONEBIT_NPU_LAX_REF} elf) + add_test(NAME npu_lax_e2e_classic + COMMAND ${CMAKE_SOURCE_DIR}/tests/npu_lax_cpp.sh $ + ${ONEBIT_NPU_LAX_MODEL} ${ONEBIT_NPU_LAX_KERNELS} ${ONEBIT_NPU_LAX_REF} classic) + endif() + if(ONEBIT_NPU_LAX_MODEL AND ONEBIT_NPU_LAX_KERNELS) + # A DeltaNet state snapshot restored after another continuation gives the logits of the + # same prefix fed from scratch, bit for bit; a Session follow-up from it, the same tokens. + add_test(NAME npu_lax_snapshot + COMMAND onebit npu-lax --model ${ONEBIT_NPU_LAX_MODEL} --kernels ${ONEBIT_NPU_LAX_KERNELS} + --snapshot-check) + # `1bit serve --device npu` on the 35B: the route, one chat round trip, no xclbin opened. + add_test(NAME npu_lax_serve + COMMAND ${CMAKE_SOURCE_DIR}/tests/npu_lax_serve.sh $ + ${ONEBIT_NPU_LAX_MODEL} ${ONEBIT_NPU_LAX_KERNELS}) + endif() endif() diff --git a/NOTICE b/NOTICE index ccf1e5d7..4a3f5397 100644 --- a/NOTICE +++ b/NOTICE @@ -68,6 +68,18 @@ ZINC (third_party/zinc) Copyright (c) 2025 ZINC Contributors License: MIT +OpenFlowLM-Next open kernels (third_party/OpenFlowLM-Next; only open_kernels/ is built) + https://github.com/bong-water-water-bong/OpenFlowLM-Next + (a fork of https://github.com/Atomic-Germ/OpenFlowLM-Next) + open_kernels/: Copyright (c) 2026 Cyrus Attoun and phlegm contributors; + portions derived from OpenFlowLM (https://github.com/OpenFlowLM/OpenFlowLM), + Copyright (c) 2026 Advanced Micro Devices, Inc. + License: MIT (open_kernels/LICENSE); the rest of the tree: Copyright (c) + OpenFlowLM Community, MIT (LICENSE_OPEN_RUNTIME.md) + Includes GEMV tile arithmetic and a smoke test adapted from vegah/LLMNpuTest + (https://github.com/vegah/LLMNpuTest), Apache-2.0 + (open_kernels/designs/rot13/LICENSE.LLMNpuTest; attributed in the file headers) + Hugging Face tokenizers (third_party/tokenizers) https://github.com/huggingface/tokenizers Copyright Hugging Face, Inc. and the tokenizers authors diff --git a/README.md b/README.md index aafed132..e40ffa74 100644 --- a/README.md +++ b/README.md @@ -40,7 +40,9 @@ serves each model behind an OpenAI-compatible API (`1bit serve`), whatever devic > vendors Lemonade: Lemonade is the host, 1bit is the engine inside it. Ported so far: HRX on AMD's live ggml-hrx > ([docs/hrx.md](docs/hrx.md)), Vulkan from upstream llama.cpp's latest release ([docs/vulkan.md](docs/vulkan.md)), the NPU engine on full ELFs with the > upstream XDNA stack pinned ([docs/npu.md](docs/npu.md); its layer kernel is not yet built from -> source), ZINC ([docs/zinc.md](docs/zinc.md)) and MLX ([docs/apple.md](docs/apple.md)). Next is step 4, +> source), ZINC ([docs/zinc.md](docs/zinc.md)) and MLX ([docs/apple.md](docs/apple.md)). Experimental, +> and served by `1bit serve --device npu` on full ELFs: Qwen3.6-35B-A3B on the NPU, the whole 40-layer +> MoE token as one runlist submit, parity against fp64, 16.3-16.5 tok/s decode ([docs/npu-lax.md](docs/npu-lax.md)). Next is step 4, > the Laya router. The working engine is being ported from > [1bit-MONSTER](https://github.com/1bit-MONSTER/1bit-MONSTER); see [docs/PORTING.md](docs/PORTING.md). > Measured results are on the [wiki](https://github.com/1bit-MONSTER/engine/wiki). This repository diff --git a/app/main.cpp b/app/main.cpp index 90412aee..3ec1d9b6 100644 --- a/app/main.cpp +++ b/app/main.cpp @@ -21,12 +21,15 @@ // OpenAI endpoints; serve's NPU route. // 1bit npu-run [options] token ids in, token ids out, on the NPU fast // lane (docs/npu.md); for checks and benchmarks. +// 1bit npu-lax [options] Qwen3.6-35B-A3B on the NPU lax kernels: greedy +// chat, or the parity check (docs/npu-lax.md). // // See docs/PORTING.md for what lands next. #ifdef ONEBIT_NPU #include "generate.h" #include "model.h" +#include "npu_lax.h" #include "unified.h" #endif #include "serve.h" @@ -56,6 +59,7 @@ void usage(FILE* out) { #ifdef ONEBIT_NPU " unified -m serve one NPU model (OpenAI endpoints)\n" " npu-run [options] generate on the NPU fast lane (npu-run --help)\n" + " npu-lax [options] Qwen3.6-35B-A3B on the NPU lax kernels (npu-lax --help)\n" #endif " version print the version\n" " help show this help\n"); @@ -148,6 +152,14 @@ int main(int argc, char** argv) { return 1; } } + if (cmd == "npu-lax") { + try { + return onebit::run_npu_lax(argc - 2, argv + 2); + } catch (const std::exception& e) { + std::fprintf(stderr, "1bit npu-lax: %s\n", e.what()); + return 1; + } + } #endif if (cmd == "version" || cmd == "--version") { std::printf("1bit %s\n", kVersion); diff --git a/app/npu_lax.cpp b/app/npu_lax.cpp new file mode 100644 index 00000000..5c750180 --- /dev/null +++ b/app/npu_lax.cpp @@ -0,0 +1,357 @@ +// 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 "npu_lax.h" + +#include "lax.h" +#include "model.h" +#include "tokenizer.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace onebit { + +namespace { + +namespace lax = npu::lax; +using clk = std::chrono::steady_clock; +double sec_since(clk::time_point t) { return std::chrono::duration(clk::now() - t).count(); } + +std::vector read_f32(const std::string& path) { + std::ifstream f(path, std::ios::binary | std::ios::ate); + if (!f) throw std::runtime_error("cannot read " + path); + std::vector v(size_t(f.tellg()) / 4); + f.seekg(0); + f.read(reinterpret_cast(v.data()), std::streamsize(v.size() * 4)); + return v; +} + +std::string sfx(int t) { return t ? "_t" + std::to_string(t) : ""; } + +// The reference's metric (open_kernels model/compare_decode.py): Pearson correlation over +// the common length, argmax (first maximum) and top-5; PASS needs finite logits, corr > +// 0.9999 and the same argmax. +bool compare(int t, const std::vector& ours, const std::vector& ref) { + const size_t n = std::min(ours.size(), ref.size()); + double ma = 0, mb = 0; + bool finite = true; + for (size_t i = 0; i < n; ++i) { + finite = finite && std::isfinite(ours[i]); + ma += ours[i]; + mb += ref[i]; + } + ma /= double(n); + mb /= double(n); + double sab = 0, saa = 0, sbb = 0; + for (size_t i = 0; i < n; ++i) { + const double a = double(ours[i]) - ma, b = double(ref[i]) - mb; + sab += a * b; + saa += a * a; + sbb += b * b; + } + const double corr = sab / std::sqrt(saa * sbb); + auto top = [n](const std::vector& v, int k) { + std::vector idx(n); + std::iota(idx.begin(), idx.end(), 0); + std::partial_sort(idx.begin(), idx.begin() + k, idx.end(), [&](int a, int b) { return v[size_t(a)] > v[size_t(b)] || (v[size_t(a)] == v[size_t(b)] && a < b); }); + idx.resize(size_t(k)); + return idx; + }; + const auto to = top(ours, 5), tr = top(ref, 5); + auto list = [](const std::vector& v) { + std::string s; + for (int x : v) s += (s.empty() ? "" : " ") + std::to_string(x); + return s; + }; + const bool ok = finite && corr > 0.9999 && to[0] == tr[0]; + std::printf("position %d: logits corr %.6f argmax ours %d ref %d top5 ours [%s] ref [%s] %s\n", t, corr, to[0], tr[0], + list(to).c_str(), list(tr).c_str(), ok ? "PASS" : "FAIL"); + return ok; +} + +int single_id(const npu::Tokenizer& tok, const char* text) { + const auto ids = tok.encode(text); + if (ids.size() != 1) throw std::runtime_error(std::string("the tokenizer has no single id for ") + text); + return ids[0]; +} + +// --snapshot-check (docs/npu-lax.md, "Chat follow-ups"): prefix P, continuation X, and a +// different continuation Y. From scratch: P + X, with X's logits. Restored: P, save the +// DeltaNet state, Y (which rewrites the state and the KV rows at X's positions), load the +// state, X again at the same positions. X's logits must be the same bytes. Then the same +// through lax::Session: P + Y generated, then P + X from the snapshot taken before Y, +// which must generate what P + X generates from scratch. +int snapshot_check_mode(lax::Decoder& dec, const std::string& model_dir, const std::string& dump) { + const npu::Tokenizer tok(model_dir + "/tokenizer.json"); + const auto P = tok.encode("<|im_start|>user\nWhat is the capital of France? Answer in one sentence.<|im_end|>\n"); + const auto X = tok.encode("<|im_start|>assistant\nThe capital of France is Paris.<|im_end|>\n<|im_start|>user\n" + "And of Germany?<|im_end|>\n<|im_start|>assistant\n\n\n\n\n"); + const auto Y = tok.encode("<|im_start|>assistant\n\nThe user asks about France, whose capital is Paris.\n" + "\n\nParis.<|im_end|>\n"); + std::printf("P %zu tokens, X %zu, Y %zu; state %zu bytes\n", P.size(), X.size(), Y.size(), dec.state_bytes()); + std::vector lg; + auto feed = [&](const std::vector& t, size_t pos, std::vector* out) { + for (size_t i = 0; i < t.size(); ++i) { + dec.feed(t[i]); + dec.run(pos + i, out != nullptr); + if (out) { + dec.logits(lg); + out->insert(out->end(), lg.begin(), lg.end()); + } + } + }; + std::vector scratch, restored; + dec.reset(); + feed(P, 0, nullptr); + feed(X, P.size(), &scratch); + + dec.reset(); + feed(P, 0, nullptr); + std::vector state(dec.state_bytes()); + auto t0 = clk::now(); + dec.save_state(state.data()); + const double save_ms = sec_since(t0) * 1000; + feed(Y, P.size(), nullptr); + t0 = clk::now(); + dec.load_state(state.data()); + const double load_ms = sec_since(t0) * 1000; + feed(X, P.size(), &restored); + std::printf("save_state %.1f ms, load_state %.1f ms\n", save_ms, load_ms); + if (!dump.empty()) { + for (auto [name, v] : {std::pair{"scratch", &scratch}, std::pair{"restored", &restored}}) + std::ofstream(dump + "/" + name + ".bin", std::ios::binary) + .write(reinterpret_cast(v->data()), std::streamsize(v->size() * 4)); + } + const size_t row = scratch.size() / X.size(); + bool all = scratch.size() == restored.size(); + size_t same = 0; + for (size_t i = 0; all && i < X.size(); ++i) + same += std::memcmp(scratch.data() + i * row, restored.data() + i * row, row * 4) == 0; + all = all && same == X.size(); + std::printf("restored snapshot: logits at X's %zu positions bit-identical to P + X from scratch: %zu of %zu %s\n", + X.size(), same, X.size(), all ? "PASS" : "FAIL"); + + // The same through Session, as serve drives it: a turn token makes the snapshot at P. + const int vocab = tok.size(); + lax::GenerateOptions opt; + opt.max_tokens = 12; + opt.stop = {single_id(tok, "<|im_end|>"), single_id(tok, "<|endoftext|>")}; + auto cat = [](std::vector a, const std::vector& b) { + a.insert(a.end(), b.begin(), b.end()); + return a; + }; + lax::Session plain(dec, vocab, {.snapshots = 0}); + const auto want = plain.generate(cat(P, X), opt); + lax::Session s(dec, vocab, {.snapshots = 2, .turn_tokens = {single_id(tok, "<|im_start|>")}}); + s.generate(cat(P, Y), opt); + const auto got = s.generate(cat(P, X), opt); + const bool ok = got.tokens == want.tokens && got.source == lax::Plan::Snapshot && got.reused == P.size(); + std::printf("session: P + X after P + Y reused %zu tokens (%s, restore %.1f ms); generated %s the %zu tokens from " + "scratch (\"%s\") %s\n", + got.reused, got.source == lax::Plan::Snapshot ? "snapshot" : "not a snapshot", got.restore_ms, + got.tokens == want.tokens ? "=" : "!=", want.tokens.size(), tok.decode(want.tokens).c_str(), + ok ? "PASS" : "FAIL"); + all = all && ok; + std::printf("%s\n", all ? "PASS" : "FAIL"); + return all ? 0 : 1; +} + +} // namespace + +int run_npu_lax(int argc, char** argv) { + const auto t_start = clk::now(); + std::string model_dir, kernel_dir, parity, dump; + int max_new = 256, tokens = 3, threads = 0, bench = 0; + bool snapshot_check = false; + lax::Transport transport = lax::Transport::Elf; + std::vector prompts; + for (int i = 0; i < argc; ++i) { + const std::string a = argv[i]; + auto next = [&]() -> std::string { + if (i + 1 >= argc) throw std::runtime_error(a + " needs a value"); + return argv[++i]; + }; + if (a == "--model") model_dir = next(); + else if (a == "--kernels") kernel_dir = next(); + else if (a == "-n" || a == "--max-new") max_new = std::stoi(next()); + else if (a == "--threads") threads = std::stoi(next()); + else if (a == "--parity") parity = next(); + else if (a == "--tokens") tokens = std::stoi(next()); + else if (a == "--dump") dump = next(); + else if (a == "--transport") transport = lax::parse_transport(next()); + else if (a == "--bench") bench = std::stoi(next()); + else if (a == "--snapshot-check") snapshot_check = true; + else if (a == "--help" || a == "-h") { + std::printf( + "usage: 1bit npu-lax --model --kernels [-n 256] [--threads N] [\"prompt\" ...]\n" + " 1bit npu-lax --model --kernels --parity [--tokens 3] [--dump ]\n" + " 1bit npu-lax --model --kernels --bench N\n" + " 1bit npu-lax --model --kernels --snapshot-check [--dump ]\n" + " --model Qwen3.6-35B-A3B Q4NX directory (model.q4nx, config.json, tokenizer.json)\n" + " --kernels scripts/build-lax.sh's /kernels (lax_l, lax_a, ln, lm_head_q8)\n" + " --transport elf (full ELFs, the default) or classic (xclbin + insts.bin)\n" + " --parity make_decode.py --requant output: runs xres.bin at position t and scores\n" + " the logits against ref_logits<_tN>.bin (corr > 0.9999, same argmax), then\n" + " again after reset(), which must give the same logits bit for bit\n" + " --bench N tokens at positions 0..N-1, twice (new and cached position configs);\n" + " prints the mean host prep, 40-layer and head times\n" + " --snapshot-check a DeltaNet state snapshot restored after another continuation must\n" + " give the logits of the same prefix fed from scratch, bit for bit; and a\n" + " Session follow-up from a snapshot the tokens of one from scratch\n" + " (--dump: scratch.bin / restored.bin, the continuation's logits, for cmp)\n" + "Without prompts, reads one prompt per line from stdin; the conversation is one session.\n"); + return 0; + } else if (!a.empty() && a[0] == '-') throw std::runtime_error("unknown option " + a); + else prompts.push_back(a); + } + if (model_dir.empty() || kernel_dir.empty()) throw std::runtime_error("--model and --kernels are required"); + + const npu::Model model(model_dir); + const auto cfg = lax::Config::from_model(model, model_dir); + std::fprintf(stderr, "[packing %s onto the NPU...]\n", model_dir.c_str()); + lax::Decoder dec(model, cfg, kernel_dir, transport, threads); + const auto& ls = dec.load_stats(); + std::fprintf(stderr, "[ready in %.1f s since process start: buffers %.1f s, packing %.2f GB %.1f s; %s]\n", + sec_since(t_start), ls.buffers_ms / 1000, double(ls.packed_bytes) / 1e9, ls.pack_ms / 1000, + dec.kernels().c_str()); + + if (!parity.empty()) { + bool all = true; + std::vector lg; + std::vector> first; + for (int t = 0; t < tokens; ++t) { + const auto x = read_f32(parity + "/xres" + std::to_string(t) + ".bin"); + if (x.size() < lax::kHidden) throw std::runtime_error("xres" + std::to_string(t) + ".bin is short"); + dec.set_residual(x.data()); + const auto t0 = clk::now(); + dec.run(size_t(t), true); + const double ms = sec_since(t0) * 1000; + dec.logits(lg); + if (!dump.empty()) + std::ofstream(dump + "/y_logits" + sfx(t) + ".bin", std::ios::binary) + .write(reinterpret_cast(lg.data()), std::streamsize(lg.size() * 4)); + std::fprintf(stderr, "[position %d: %.1f ms]\n", t, ms); + all = compare(t, lg, read_f32(parity + "/ref_logits" + sfx(t) + ".bin")) && all; + first.push_back(lg); + } + // A second sequence after reset() must see nothing of the first. + dec.reset(); + for (int t = 0; t < tokens; ++t) { + const auto x = read_f32(parity + "/xres" + std::to_string(t) + ".bin"); + dec.set_residual(x.data()); + dec.run(size_t(t), true); + dec.logits(lg); + const bool same = lg == first[size_t(t)]; + std::printf("position %d after reset: logits %s\n", t, same ? "bit-identical PASS" : "differ FAIL"); + all = same && all; + } + std::printf("%s\n", all ? "PASS" : "FAIL"); + return all ? 0 : 1; + } + + if (snapshot_check) return snapshot_check_mode(dec, model_dir, dump); + + if (bench > 0) { + const npu::Tokenizer tok(model_dir + "/tokenizer.json"); + const int token = single_id(tok, "<|im_start|>"); + for (int pass = 0; pass < 2; ++pass) { + dec.reset(); + lax::RunTimes sum; + const auto t0 = clk::now(); + for (int p = 0; p < bench; ++p) { + dec.feed(token); + dec.run(size_t(p), true); + const auto& r = dec.last_run(); + sum.prep_ms += r.prep_ms; + sum.layers_ms += r.layers_ms; + sum.head_ms += r.head_ms; + sum.ahead_ms += r.ahead_ms; + } + const double total = sec_since(t0) * 1000 / bench; + std::printf("bench %s: %d tokens, per token %.2f ms (%.1f tok/s): prep %.3f ms, 40 layers %.2f ms " + "(%.3f ms of it preparing the next), norm + lm head %.2f ms\n", + pass ? "cached positions" : "new positions", bench, total, 1000 / total, sum.prep_ms / bench, + sum.layers_ms / bench, sum.ahead_ms / bench, sum.head_ms / bench); + } + std::fprintf(stderr, "[%s]\n", dec.kernels().c_str()); + return 0; + } + + const npu::Tokenizer tok(model_dir + "/tokenizer.json"); + const int vocab = tok.size(); // the lm head's rows past the tokenizer are padding + const int im_end = single_id(tok, "<|im_end|>"), eot = single_id(tok, "<|endoftext|>"); + size_t pos = 0; + auto turn = [&](const std::string& text) { + // Qwen3.6's chat template, thinking off: an empty think block opens the answer. + std::vector ids = tok.encode("<|im_start|>user\n" + text + "<|im_end|>\n<|im_start|>assistant\n\n\n\n\n"); + if (pos) { + auto nl = tok.encode("\n"); // after the previous turn's <|im_end|> + ids.insert(ids.begin(), nl.begin(), nl.end()); + } + if (pos + ids.size() + size_t(max_new) + 1 > size_t(lax::kMaxContext)) { + std::fprintf(stderr, "[context full (%zu of %d positions)]\n", pos, lax::kMaxContext); + return false; + } + const auto t1 = clk::now(); + for (size_t i = 0; i < ids.size(); ++i) { + dec.feed(ids[i]); + dec.run(pos++, i + 1 == ids.size()); // the head only where a next token is wanted + } + int next = dec.argmax(vocab); + const double prompt_s = sec_since(t1); + const auto t2 = clk::now(); + std::vector out; + std::string shown; + while (next != im_end && next != eot && int(out.size()) < max_new) { + out.push_back(next); + const std::string s = tok.decode(out); + std::fwrite(s.data() + shown.size(), 1, s.size() - shown.size(), stdout); + std::fflush(stdout); + shown = s; + dec.feed(next); + dec.run(pos++, true); + next = dec.argmax(vocab); + } + const double gen_s = sec_since(t2); + dec.feed(next); // the end-of-turn token enters the cache too + dec.run(pos++, false); + std::printf("\n"); + std::fflush(stdout); + std::fprintf(stderr, "[prompt %zu tok in %.2f s (%.1f tok/s), %zu tok in %.2f s = %.1f tok/s, context %zu]\n", + ids.size(), prompt_s, double(ids.size()) / prompt_s, out.size(), gen_s, + double(out.size()) / std::max(gen_s, 1e-9), pos); + return true; + }; + if (!prompts.empty()) { + for (const auto& p : prompts) + if (!turn(p)) break; + } else { + for (std::string line; std::getline(std::cin, line);) + if (line.find_first_not_of(" \t\r") != std::string::npos && !turn(line)) break; + } + return 0; +} + +} // namespace onebit diff --git a/app/npu_lax.h b/app/npu_lax.h new file mode 100644 index 00000000..1d5f3a50 --- /dev/null +++ b/app/npu_lax.h @@ -0,0 +1,24 @@ +// 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. + +// `1bit npu-lax`: Qwen3.6-35B-A3B on the NPU's lax kernels (docs/npu-lax.md), greedy chat +// or the three-position parity check against the fp64 reference. +#pragma once + +namespace onebit { + +int run_npu_lax(int argc, char** argv); + +} // namespace onebit diff --git a/app/serve.cpp b/app/serve.cpp index d14164a7..8fd914d2 100644 --- a/app/serve.cpp +++ b/app/serve.cpp @@ -14,6 +14,7 @@ // limitations under the License. // `1bit serve -m [--port 8000] [--device auto|npu|vulkan|hrx|rocm|zinc] [--lean] [--mtp HEAD] [--adaptive]` +// [--npu-kernels DIR] [--npu-transport elf|classic] [--npu-snapshots N] // // The engine's one front door (docs/serve.md): one model per process behind an // OpenAI-compatible API. It is how the engine runs inside Lemonade: Lemonade's @@ -26,6 +27,8 @@ // // Which device runs the model follows from the model and --device: // an NPU model directory (model.q4nx + npu/) -> the NPU fast lane, in process +// Qwen3.6-35B-A3B's Q4NX directory -> the NPU lax decode from full ELFs, in +// process (docs/npu-lax.md) // a .gguf, --device vulkan -> the upstream llama.cpp build's // llama-server (else the HRX build's) // a .gguf, --device hrx -> the HRX build's llama-server @@ -103,6 +106,8 @@ struct Options { bool adaptive = false; int adaptive_at = 1; std::string embed, rerank; // RAG: an embedding model and a reranker, served beside the chat model + std::string npu_kernels, npu_transport; // the lax decode's kernels and their transport + std::string npu_snapshots; // the lax decode's state snapshots kept }; // HRX dlopens the HSA runtime, and a distro libhsa rejects gfx1151's @@ -581,8 +586,10 @@ void usage(FILE* out) { " [--parallel N] N requests decoded together (continuous batching)\n" " [--adaptive] [--adaptive-at N] Vulkan (+MTP) for N in flight (default 1), ROCm batches the rest\n" " [--embed MODEL.gguf] [--rerank MODEL.gguf] RAG: /v1/embeddings and /v1/rerank\n" - " : an NPU model directory (model.q4nx + npu/), a .gguf file, or with\n" - " --device mlx a Hugging Face id (mlx-community/...)\n"); + " [--npu-kernels DIR] [--npu-transport elf|classic] Qwen3.6-35B-A3B's lax kernels\n" + " [--npu-snapshots 4] its DeltaNet state snapshots for chat follow-ups (70 MB each)\n" + " : an NPU model directory (model.q4nx + npu/, or Qwen3.6-35B-A3B's with the lax\n" + " kernels), a .gguf file, or with --device mlx a Hugging Face id (mlx-community/...)\n"); } } // namespace @@ -616,6 +623,9 @@ int run_serve(int argc, char** argv) { else if (a == "--adaptive-at") o.adaptive_at = std::stoi(next()); else if (a == "--embed") o.embed = next(); else if (a == "--rerank") o.rerank = next(); + else if (a == "--npu-kernels") o.npu_kernels = next(); + else if (a == "--npu-transport") o.npu_transport = next(); + else if (a == "--npu-snapshots") o.npu_snapshots = next(); else if (a == "-h" || a == "--help") { usage(stdout); return 0; } else throw std::runtime_error("unknown option " + a); } @@ -626,10 +636,16 @@ int run_serve(int argc, char** argv) { if (o.device != "auto" && o.device != "npu") throw std::runtime_error("an NPU model directory runs on --device npu"); #ifdef ONEBIT_NPU - if (!is_npu_model_dir(o.model)) throw std::runtime_error(o.model + " is not an NPU model directory"); - // The NPU lane serves in process (unified.cpp). + if (!is_npu_model_dir(o.model, o.npu_kernels)) + throw std::runtime_error(o.model + " is not an NPU model directory (the fast lane needs npu/; " + "Qwen3.6-35B-A3B needs the lax kernels: --npu-kernels, /npu/lax or " + "$ONEBIT_NPU_LAX_KERNELS)"); + // The NPU serves in process (unified.cpp). std::vector args = {"-m", o.model, "-p", std::to_string(o.port), "--host", o.host}; if (!o.alias.empty()) { args.push_back("--alias"); args.push_back(o.alias); } + if (!o.npu_kernels.empty()) { args.push_back("--kernels"); args.push_back(o.npu_kernels); } + if (!o.npu_transport.empty()) { args.push_back("--transport"); args.push_back(o.npu_transport); } + if (!o.npu_snapshots.empty()) { args.push_back("--snapshots"); args.push_back(o.npu_snapshots); } std::vector av; for (auto& s : args) av.push_back(s.data()); return run_unified(int(av.size()), av.data()); diff --git a/app/unified.cpp b/app/unified.cpp index 509cf78c..b66f1cf9 100644 --- a/app/unified.cpp +++ b/app/unified.cpp @@ -13,19 +13,25 @@ // See the License for the specific language governing permissions and // limitations under the License. -// `1bit unified -m -p `: one native model on the NPU fast lane, -// behind OpenAI endpoints. It is `1bit serve`'s NPU route (serve.cpp): +// `1bit unified -m -p `: one native model on the NPU, behind OpenAI +// endpoints. It is `1bit serve`'s NPU route (serve.cpp). The model decides the path +// (npu_model_kind): Qwen2/Qwen3 dense models run on the fast lane (docs/npu.md), +// Qwen3.6-35B-A3B on the lax decode from full ELFs (docs/npu-lax.md). // // GET /health, /v1/health 200 once the model is on the device // GET /v1/models the one model // POST /v1/chat/completions ChatML prompt, greedy decoding, optional SSE stream // POST /v1/completions raw prompt // -// Requests run one at a time: the lane holds one KV cache, and each request -// rewrites it from position 0. +// Requests run one at a time: the device holds one cache. The fast lane rewrites it from +// position 0 per request. The lax decode continues it when a request extends the last +// one's tokens exactly, or restores a snapshot of the DeltaNet state taken at an earlier +// prompt's prefix (--snapshots, npu/lax_turns.h): a chat follow-up re-feeds only the turns +// after the one its history last diverged in. It starts over otherwise. #include "unified.h" #include "generate.h" +#include "lax.h" #include "model.h" #include "tokenizer.h" @@ -35,11 +41,14 @@ #include #include #include +#include #include +#include #include #include #include #include +#include namespace onebit { @@ -51,14 +60,53 @@ struct Engine { std::string id; std::unique_ptr model; std::unique_ptr tok; - std::unique_ptr lane; + std::unique_ptr lane; // the fast lane, or + std::unique_ptr lax_cfg; + std::unique_ptr lax; // the lax decode and its conversation + std::unique_ptr session; + std::vector lax_stop; + int max_context = npu::Lane::kMaxContext; + float default_penalty = 1.1f; // greedy 0.6B models loop without it; the 35B does not need it + bool open_think = false; // Qwen3.6's template opens "\n" when thinking is on std::mutex mu; + + // reused: prompt tokens already on the device (the lax decode's conversation); cache: + // where from ("live", "snapshot"), empty when none. + npu::GenerateResult generate(const std::vector& ids, const npu::GenerateOptions& opt, size_t& reused, + std::string& cache) { + reused = 0; + cache.clear(); + if (lane) return npu::generate(*lane, *model, ids, opt); + npu::lax::GenerateOptions lo; + lo.max_tokens = opt.max_tokens; + lo.repetition_penalty = opt.repetition_penalty; + lo.penalty_window = opt.penalty_window; + lo.stop = lax_stop; + lo.on_token = opt.on_token; + const auto r = session->generate(ids, lo); + npu::GenerateResult out; + out.tokens = r.tokens; + out.prefill_ms = r.prefill_ms; + out.decode_ms = r.decode_ms; + out.stopped_at_eos = r.stopped_at_eos; + reused = r.reused; + if (r.source != npu::lax::Plan::Scratch) cache = r.source == npu::lax::Plan::Live ? "live" : "snapshot"; + if (r.source == npu::lax::Plan::Snapshot || r.snapshots_taken) + std::fprintf(stderr, "1bit unified: %zu of %zu prompt tokens from %s (restore %.1f ms), %d snapshot%s taken " + "(%.1f ms), %zu kept (%.0f MB)\n", + r.reused, ids.size(), cache.empty() ? "nothing" : cache.c_str(), r.restore_ms, r.snapshots_taken, + r.snapshots_taken == 1 ? "" : "s", r.snapshot_ms, session->snapshots().size(), + double(session->snapshots().bytes()) / 1e6); + return out; + } }; // The chat scaffold for the model families the lane serves. Hugging Face ships // the template as Jinja; ChatML is what Qwen2 and Qwen3 render it to for plain // text messages, with generation starting after "<|im_start|>assistant\n". -std::string chatml(const json& messages, bool thinking = true) { +// open_think: the template opens the think block itself when thinking is on (Qwen3.6's +// does; Qwen3's leaves it to the model). +std::string chatml(const json& messages, bool thinking = true, bool open_think = false) { std::string p; for (const auto& m : messages) { std::string content; @@ -75,6 +123,7 @@ std::string chatml(const json& messages, bool thinking = true) { // chat_template_kwargs.enable_thinking = false: what Qwen3's own template // emits, an empty think block, so the model answers directly. if (!thinking) p += "\n\n\n\n"; + else if (open_think) p += "\n"; return p; } @@ -108,7 +157,7 @@ void handle_generate(Engine& e, const httplib::Request& req, httplib::Response& bool thinking = true; if (body.contains("chat_template_kwargs") && body["chat_template_kwargs"].is_object()) thinking = body["chat_template_kwargs"].value("enable_thinking", true); - prompt = chat ? chatml(body.at("messages"), thinking) : body.at("prompt").get(); + prompt = chat ? chatml(body.at("messages"), thinking, e.open_think) : body.at("prompt").get(); } catch (const std::exception& ex) { res.status = 400; res.set_content(json{{"error", {{"message", std::string("bad request: ") + ex.what()}}}}.dump(), "application/json"); @@ -116,8 +165,11 @@ void handle_generate(Engine& e, const httplib::Request& req, httplib::Response& } npu::GenerateOptions opt; opt.max_tokens = body.value("max_completion_tokens", body.value("max_tokens", 1024)); - opt.repetition_penalty = body.value("repetition_penalty", 1.1f); + opt.repetition_penalty = body.value("repetition_penalty", e.default_penalty); const bool stream = body.value("stream", false); + // stream_options.include_usage: a last chunk with no choices carries the usage, as OpenAI's. + const bool stream_usage = stream && body.contains("stream_options") && body["stream_options"].is_object() && + body["stream_options"].value("include_usage", false); const std::string id = now_id(chat ? "chatcmpl-" : "cmpl-"); const long created = long(std::chrono::duration_cast( std::chrono::system_clock::now().time_since_epoch()).count()); @@ -127,7 +179,7 @@ void handle_generate(Engine& e, const httplib::Request& req, httplib::Response& std::lock_guard lock(e.mu); const auto ids = e.tok->encode(prompt); if (ids.empty()) throw std::runtime_error("empty prompt"); - const int room = npu::Lane::kMaxContext - int(ids.size()); + const int room = e.max_context - int(ids.size()); if (room < 1) throw std::runtime_error("prompt longer than the context"); opt.max_tokens = std::min(opt.max_tokens, room); std::string pending; @@ -139,9 +191,11 @@ void handle_generate(Engine& e, const httplib::Request& req, httplib::Response& pending.erase(0, n); return on_text(out); }; - auto r = npu::generate(*e.lane, *e.model, ids, opt); + size_t reused = 0; + std::string cache; + auto r = e.generate(ids, opt, reused, cache); if (!pending.empty()) on_text(pending); - return std::make_pair(r, ids.size()); + return std::make_tuple(r, ids.size(), reused, cache); }; // By value: for a stream, httplib calls the content provider after this // handler has returned. @@ -152,14 +206,38 @@ void handle_generate(Engine& e, const httplib::Request& req, httplib::Response& return "data: " + c.dump() + "\n\n"; }; + auto usage = [](const npu::GenerateResult& r, size_t n_prompt, size_t n_cached) { + const size_t completion = r.tokens.size() - (r.stopped_at_eos ? 1 : 0); + return json{{"prompt_tokens", n_prompt}, + {"completion_tokens", completion}, + {"total_tokens", n_prompt + completion}, + {"prompt_tokens_details", {{"cached_tokens", n_cached}}}}; + }; + auto timings = [](const npu::GenerateResult& r, const std::string& cache) { + json t{{"prompt_ms", r.prefill_ms}, + {"predicted_ms", r.decode_ms}, + {"predicted_per_second", r.tokens.empty() ? 0.0 : 1000.0 * double(r.tokens.size()) / r.decode_ms}}; + if (!cache.empty()) t["cache"] = cache; + return t; + }; + if (stream) { - res.set_chunked_content_provider("text/event-stream", [run, chunk](size_t, httplib::DataSink& sink) mutable { + res.set_chunked_content_provider("text/event-stream", [run, chunk, usage, timings, stream_usage, id, object, created, + model = e.id](size_t, httplib::DataSink& sink) mutable { try { - auto [r, n_prompt] = run([&](const std::string& t) { + auto [r, n_prompt, n_cached, cache] = run([&](const std::string& t) { const std::string c = chunk(t, nullptr); return sink.write(c.data(), c.size()); }); - const std::string end = chunk("", r.stopped_at_eos ? "stop" : "length") + "data: [DONE]\n\n"; + std::string end = chunk("", r.stopped_at_eos ? "stop" : "length"); + if (stream_usage) { + const json u{{"id", id}, {"object", object}, + {"created", created}, {"model", model}, + {"choices", json::array()}, {"usage", usage(r, n_prompt, n_cached)}, + {"timings", timings(r, cache)}}; + end += "data: " + u.dump() + "\n\n"; + } + end += "data: [DONE]\n\n"; sink.write(end.data(), end.size()); } catch (const std::exception& ex) { const std::string err = "data: " + json{{"error", {{"message", ex.what()}}}}.dump() + "\n\n"; @@ -172,23 +250,18 @@ void handle_generate(Engine& e, const httplib::Request& req, httplib::Response& } try { std::string text; - auto [r, n_prompt] = run([&](const std::string& t) { + auto [r, n_prompt, n_cached, cache] = run([&](const std::string& t) { text += t; return true; }); - const size_t completion = r.tokens.size() - (r.stopped_at_eos ? 1 : 0); json out{{"id", id}, {"object", object}, {"created", created}, {"model", e.id}}; const char* finish = r.stopped_at_eos ? "stop" : "length"; out["choices"] = json::array({chat ? json{{"index", 0}, {"message", {{"role", "assistant"}, {"content", text}}}, {"finish_reason", finish}} : json{{"index", 0}, {"text", text}, {"finish_reason", finish}}}); - out["usage"] = {{"prompt_tokens", n_prompt}, - {"completion_tokens", completion}, - {"total_tokens", n_prompt + completion}}; - out["timings"] = {{"prompt_ms", r.prefill_ms}, - {"predicted_ms", r.decode_ms}, - {"predicted_per_second", r.tokens.empty() ? 0.0 : 1000.0 * double(r.tokens.size()) / r.decode_ms}}; + out["usage"] = usage(r, n_prompt, n_cached); + out["timings"] = timings(r, cache); res.set_content(out.dump(), "application/json"); } catch (const std::exception& ex) { res.status = 500; @@ -200,17 +273,55 @@ void handle_generate(Engine& e, const httplib::Request& req, httplib::Response& std::string npu_kernel_dir(const std::string& model_dir) { return model_dir + "/npu"; } -bool is_npu_model_dir(const std::string& dir) { +namespace { + +std::string model_type(const std::string& dir) { + std::ifstream f(dir + "/config.json"); + if (!f) return ""; + try { + const auto j = json::parse(f); + return j.value("model_type", ""); + } catch (const std::exception&) { + return ""; + } +} + +bool has_lax_kernels(const std::string& k) { namespace fs = std::filesystem; - for (const char* f : {"model.q4nx", "config.json", "tokenizer.json", "npu/layer_ctx1.elf", "npu/layer_ctx2.elf", - "npu/layer_ctx17.elf", "npu/lmhead.elf", "npu/layer.pdi"}) - if (!fs::exists(fs::path(dir) / f)) return false; + for (const char* kind : {"lax_l", "lax_a", "ln", "lm_head_q8"}) + if (!fs::exists(fs::path(k) / kind / "insts.elf")) return false; return true; } +} // namespace + +std::string lax_kernel_dir(const std::string& dir, const std::string& lax_kernels) { + if (!lax_kernels.empty()) return lax_kernels; + if (has_lax_kernels(dir + "/npu/lax")) return dir + "/npu/lax"; + if (const char* env = std::getenv("ONEBIT_NPU_LAX_KERNELS"); env && *env) return env; + return ""; +} + +NpuModel npu_model_kind(const std::string& dir, const std::string& lax_kernels) { + namespace fs = std::filesystem; + for (const char* f : {"model.q4nx", "config.json", "tokenizer.json"}) + if (!fs::exists(fs::path(dir) / f)) return NpuModel::None; + if (model_type(dir) == "qwen3_5_moe") { + const std::string k = lax_kernel_dir(dir, lax_kernels); + return !k.empty() && has_lax_kernels(k) ? NpuModel::Lax : NpuModel::None; + } + for (const char* f : {"npu/layer_ctx1.elf", "npu/layer_ctx2.elf", "npu/layer_ctx17.elf", "npu/lmhead.elf", "npu/layer.pdi"}) + if (!fs::exists(fs::path(dir) / f)) return NpuModel::None; + return NpuModel::Lane; +} + +bool is_npu_model_dir(const std::string& dir, const std::string& lax_kernels) { + return npu_model_kind(dir, lax_kernels) != NpuModel::None; +} + int run_unified(int argc, char** argv) { - std::string model_dir, host = "127.0.0.1", alias; - int port = 8000; + std::string model_dir, host = "127.0.0.1", alias, kernels, transport = "elf"; + int port = 8000, snapshots = 4; for (int i = 0; i < argc; ++i) { const std::string a = argv[i]; auto next = [&]() -> std::string { @@ -221,13 +332,23 @@ int run_unified(int argc, char** argv) { else if (a == "-p" || a == "--port") port = std::stoi(next()); else if (a == "--host") host = next(); else if (a == "--alias") alias = next(); + else if (a == "--kernels") kernels = next(); + else if (a == "--transport") transport = next(); + else if (a == "--snapshots") snapshots = std::stoi(next()); else if (a == "--help" || a == "-h") { std::printf("usage: 1bit unified -m [-p 8000] [--host 127.0.0.1] [--alias NAME]\n" - " holds model.q4nx, config.json, tokenizer.json and npu/ (the lane's kernels)\n"); + " [--kernels DIR] [--transport elf|classic] [--snapshots 4]\n" + " holds model.q4nx, config.json, tokenizer.json, and for the fast lane\n" + " npu/ (its kernels); Qwen3.6-35B-A3B runs on the lax kernels in --kernels,\n" + " else /npu/lax, else $ONEBIT_NPU_LAX_KERNELS (docs/npu-lax.md)\n" + " --transport the lax kernels as full ELFs (default) or classic xclbin, for A/B\n" + " --snapshots DeltaNet state snapshots the lax decode keeps in host memory for chat\n" + " follow-ups (70 MB each; 0: only exact extensions reuse the cache)\n"); return 0; } else throw std::runtime_error("unknown option " + a); } if (model_dir.empty()) throw std::runtime_error("-m is required"); + if (snapshots < 0) throw std::runtime_error("--snapshots must be 0 or more"); // Lemonade hands over the checkpoint path, which may name the container file. if (std::filesystem::is_regular_file(model_dir)) model_dir = std::filesystem::path(model_dir).parent_path().string(); @@ -273,12 +394,37 @@ int run_unified(int argc, char** argv) { std::fprintf(stderr, "1bit unified: loading %s on the NPU\n", model_dir.c_str()); const auto t0 = std::chrono::steady_clock::now(); + const NpuModel kind = npu_model_kind(model_dir, kernels); + if (kind == NpuModel::None) + throw std::runtime_error(model_dir + " is not an NPU model directory (docs/serve.md)"); e.model = std::make_unique(model_dir); - const std::string type = e.model->dims().model_type; - if (type != "qwen3" && type != "qwen2") - throw std::runtime_error("model_type '" + type + "': the fast lane serves the Qwen2/Qwen3 layer kernel only"); e.tok = std::make_unique(model_dir + "/tokenizer.json"); - e.lane = std::make_unique(*e.model, npu_kernel_dir(model_dir)); + if (kind == NpuModel::Lax) { + const std::string k = lax_kernel_dir(model_dir, kernels); + e.lax_cfg = std::make_unique(npu::lax::Config::from_model(*e.model, model_dir)); + e.lax = std::make_unique(*e.model, *e.lax_cfg, k, npu::lax::parse_transport(transport)); + for (const char* s : {"<|im_end|>", "<|endoftext|>"}) { + const auto ids = e.tok->encode(s); + if (ids.size() == 1) e.lax_stop.push_back(ids[0]); + } + // Snapshots before each prompt's last <|im_start|>: a follow-up repeats every message + // before the answer it re-renders. + npu::lax::SessionOptions so; + so.snapshots = size_t(snapshots); + if (const auto ids = e.tok->encode("<|im_start|>"); ids.size() == 1) so.turn_tokens = ids; + // Rows past the tokenizer are padding. + e.session = std::make_unique(*e.lax, e.tok->size(), std::move(so)); + e.max_context = npu::lax::kMaxContext; + e.default_penalty = 1.0f; + e.open_think = true; + std::fprintf(stderr, "1bit unified: Qwen3.6-35B-A3B on the lax decode (%s); up to %d state snapshots of %.0f MB\n", + e.lax->kernels().c_str(), snapshots, double(e.lax->state_bytes()) / 1e6); + } else { + const std::string type = e.model->dims().model_type; + if (type != "qwen3" && type != "qwen2") + throw std::runtime_error("model_type '" + type + "': the fast lane serves the Qwen2/Qwen3 layer kernel only"); + e.lane = std::make_unique(*e.model, npu_kernel_dir(model_dir)); + } ready = true; std::fprintf(stderr, "1bit unified: %s ready in %.0f ms on %s:%d\n", e.id.c_str(), std::chrono::duration(std::chrono::steady_clock::now() - t0).count(), host.c_str(), diff --git a/app/unified.h b/app/unified.h index 82181789..64a1fa59 100644 --- a/app/unified.h +++ b/app/unified.h @@ -23,9 +23,15 @@ namespace onebit { // argv holds the options after "unified". int run_unified(int argc, char** argv); -// A servable NPU model directory: model.q4nx, config.json, tokenizer.json and -// npu/ with the lane's kernels (docs/npu.md). -bool is_npu_model_dir(const std::string& dir); +// What runs an NPU model directory (model.q4nx, config.json, tokenizer.json): +// Lane a Qwen2/Qwen3 dense model with npu/ holding the fast lane's kernels (docs/npu.md) +// Lax Qwen3.6-35B-A3B (model_type qwen3_5_moe) on the lax kernels (docs/npu-lax.md), +// found at lax_kernels, else /npu/lax, else $ONEBIT_NPU_LAX_KERNELS +enum class NpuModel { None, Lane, Lax }; +NpuModel npu_model_kind(const std::string& dir, const std::string& lax_kernels = ""); +bool is_npu_model_dir(const std::string& dir, const std::string& lax_kernels = ""); std::string npu_kernel_dir(const std::string& model_dir); +// The lax kernel directory for dir, as npu_model_kind looks it up; empty if none. +std::string lax_kernel_dir(const std::string& dir, const std::string& lax_kernels = ""); } // namespace onebit diff --git a/docs/PORTING.md b/docs/PORTING.md index 3de2d919..56e20ebc 100644 --- a/docs/PORTING.md +++ b/docs/PORTING.md @@ -24,6 +24,7 @@ Each step below is one PR (or a short series) that builds and runs on Strix Halo | 1 | **The engine runs inside Lemonade.** Lemonade runs `1bit serve` as a backend: one model per process behind an OpenAI-compatible API | `tools/unified_server.cpp` | **landed** ([docs/serve.md](serve.md)): NPU, Vulkan, HRX and ZINC pass `serve_e2e` on Strix Halo; `smoke_serve` runs in CI. The engine first embedded Lemonade (v11.9.0 plus local recipes); after geramyL's review it is embedded *into* Lemonade instead: the vendored copy is removed, and the Lemonade recipe that runs `1bit serve` (`onebit`) lives in our fork [1bit-MONSTER/lemonade](https://github.com/1bit-MONSTER/lemonade) ([docs/lemonade.md](lemonade.md)) | | 2 | **HRX with Vulkan.** llama.cpp with `GGML_HRX=ON` and `GGML_VULKAN=ON` in one build, on AMD's tested pair | `1bit-MONSTER/llama.cpp` `1bit/hrx-vulkan-patched` (AMD's `hrx-graph-develop-v2` plus our commits) + `ROCm/hrx-system`, pinned from `ROCm/ggml-staging-automation` | **landed** ([docs/hrx.md](hrx.md)): `1bit serve` runs the same checkpoint on `Vulkan0` and `HRX0`; kept current by `bump-hrx.yml` | | 3 | **NPU engine.** Full ELFs only, and the open 16-tile layer kernel built from source | `engine/npu` (`npu_engine_universal.cpp`, `I8Ctx::init_elf`); ELF dispatch table on `backup/iso-build-elf-native-2026-09-22`; kernel on `bench/fastlane-16tile-corrections-2026-09-22` | **3a–3c landed** ([docs/npu.md](npu.md)). 3a (#8): full ELFs generated in C++, reproducing all 5,153 captured contexts. 3b/3c (#9): the lane runtime (logits bit-identical to the reference lane, 24/24 steps, 11.0 ms/token), tokenizer, `1bit unified`, and Qwen3-0.6B served on the NPU (now through `1bit serve`). The XDNA driver and XRT are pinned upstream and built privately (#12). **Open:** the layer kernel and lm-head artifacts are not yet built from source (npu.md, "Open") | +| 3x | **35B MoE on the NPU (experimental).** The whole 40-layer Qwen3.6-35B-A3B token as one runlist submit, experts routed on the device | `third_party/OpenFlowLM-Next` `open_kernels/designs/layer_x/lax.py`, `scripts/build-lax.sh` | **landed as experimental** ([docs/npu-lax.md](npu-lax.md)): kernels built from the pinned source, `tests/npu_lax_parity.sh` passes (3 positions vs fp64, argmax 846 / 198 / 3710), 16.1 tok/s chat. C++ driver `1bit npu-lax` (no Python): packer byte-identical to the Python one on all 83 buffers, same parity, ready in 3.6-3.8 s. **Full ELFs** (the default; no xclbin opened): the seven ELFs byte-identical to the reference harness's, logits bit-identical to the classic path, 16.3-16.5 tok/s decode vs 16.0-16.1 classic; `1bit serve --device npu` routes the model there and answers "The capital of France is Paris." (`tests/npu_lax_serve.sh`). **Open:** batched prefill; memory (22 GB pinned) | | + | **Linux kernel.** The kernel that provides `amdxdna` and `amdgpu`, pinned to upstream | `torvalds/linux` release tags; config from the Strix Halo kernel of 2026-09-23 | **pinned** ([docs/kernel.md](kernel.md)): v7.3-rc4 builds into Debian packages with `amdxdna` in-tree; kept current by `bump-linux.yml`. Installing it on Strix Halo is a separate, deliberate step | | 4 | **Laya router.** A non-autoregressive scorer that picks where each request runs | `src/laya_scorer.cpp`, `include/laya_scorer.h` on `backup/laya-and-results-2026-09-22`; model at `~/models/laya` | **pinned** ([docs/laya.md](laya.md)): source `NandhaKishorM/laya` + the three HF checkpoints at one revision, hash-verified fetch, `bump-laya.yml`. Next: the C++ scorer, gated against the Python reference; then routing | | 5 | **Every HF model, kept current.** The architecture registry (569 tokens mapping 2,030 HF arch strings) and the daily HF census that finds new architectures and proposes mappings | `src/model_registry*.cpp`, `Testing/census_*.py` and `.json`, `.github/workflows/census-{watch,sweep,autopr}.yml` | the census runs daily in CI; docs report *mapped* and *run and checked* counts separately | diff --git a/docs/npu-lax.md b/docs/npu-lax.md new file mode 100644 index 00000000..7a12a3d3 --- /dev/null +++ b/docs/npu-lax.md @@ -0,0 +1,449 @@ + +# NPU: Qwen3.6-35B-A3B, one runlist submit per token (lax) + +The whole 35B MoE token runs on the NPU as **one XRT runlist submit**: + +- 40 layers: 30 DeltaNet linear-attention layers and 10 full-attention layers. +- Each layer's routed experts are chosen and fetched on the device. +- The final norm and the lm head are two more runs. + +It matches an fp64 reference, and it decodes at 16.3-16.5 tok/s. + +**Status: experimental, and served by `1bit serve --device npu` on full ELFs** +([serve.md](serve.md)). The kernels are the merged whole-layer design `lax` from the open +kernels in `third_party/OpenFlowLM-Next`. Two drivers run them: + +- The engine's C++ driver (`npu/lax*.{h,cpp}`): `1bit serve` and `1bit npu-lax`. It needs + no Python. It loads the kernels as full ELFs (PDI + control code, no xclbin, rule 4) by + default, or through XRT's classic xclbin path for A/B. See "The C++ driver" and "Full + ELFs". +- The Python driver in the pinned tree, which the C++ one is checked against. + +## The pin + +`third_party/OpenFlowLM-Next` pins branch `1bit/lax-35b` of +[bong-water-water-bong/OpenFlowLM-Next](https://github.com/bong-water-water-bong/OpenFlowLM-Next) +(MIT). That branch holds the fixes and optimisations below. Upstream's +`open_kernels/designs/layer_x/35b-whole-layer-runlist-status.md` records every +measurement behind them. + +## Build and check + +```sh +git submodule update --init third_party/OpenFlowLM-Next +IRON_VENV= AIETOOLS=/aietools \ + scripts/build-lax.sh ~/.cache/lax # ~3.5 min: 4 kernel builds + the harness + +# the fp64 reference, once (CPU, ~15 min, writes 40 x 512 MB pools into --pool-dir) +cd ~/.cache/lax/src/open_kernels +python3 model/make_decode.py --requant --model-dir --layers 40 --tokens 3 \ + --out ~/.cache/lax-ref --pool-dir ~/.cache/lax-ref/pools + +tests/npu_lax_parity.sh ~/.cache/lax ~/.cache/lax-ref # from the engine checkout +``` + +The kernels land in `~/.cache/lax/kernels` in the layout under "Kernel directory". +`` is a Q4NX model directory (`model.q4nx`, `config.json`, `tokenizer.json`). +The weights are packed from `model.q4nx` at load time and streamed into the XRT +buffers, so no pre-packed files are needed. `--requant` is required, because the `lax` +builds are all-q4_1. + +Chat: + +```sh +cd ~/.cache/lax/src/open_kernels +python3 model/lax_chat.py --model-dir --lax-l ~/.cache/lax/kernels/lax_l \ + --lax-a ~/.cache/lax/kernels/lax_a "What is the capital of France? Answer in one sentence." +``` + +Without prompts on the command line, it reads prompts from stdin. The chat is one +device session, so the KV cache and the DeltaNet state carry across turns. + +## The C++ driver + +`1bit npu-lax` does what the pinned tree's `model/lax_chat.py` and `model/lax_decode_cfg.py` +do, but in the engine: + +- It packs every buffer from `model.q4nx` in parallel threads, straight into the XRT + buffers. There is no pipe, no helper process and nothing on disk. +- It builds the 40-run runlist on one context. Each layer's arguments are `pool xres + consts kv|dkv act ptab state|dstate cfg`. +- It patches the full-attention stream for each position, then runs the norm and the lm + head and takes the greedy argmax. + +The packer is `npu/lax_pack.{h,cpp}`. It is re-implemented from the recipe's +specification (`recipes/qwen36moe.py` `pack_plan` and `recipes/pack.py`, with every +projection at q4_1 as `--requant` sets it), not copied. The position patches and the +layer config words are `npu/lax_stream.{h,cpp}`. The kernel loading sits behind +`npu/lax_kernels.h`: `elf_kernels` (the default, `npu/lax_elf*.{h,cpp}`) and +`classic_kernels` (`--transport classic`) run the same decoder. + +```sh +cmake -S . -B build -DONEBIT_NPU=ON \ + -DONEBIT_NPU_LAX_MODEL= -DONEBIT_NPU_LAX_KERNELS=~/.cache/lax/kernels \ + -DONEBIT_NPU_LAX_REF=~/.cache/lax-ref +cmake --build build +build/1bit npu-lax --model --kernels ~/.cache/lax/kernels \ + "What is the capital of France? Answer in one sentence." +build/1bit npu-lax --model --kernels ~/.cache/lax/kernels --parity ~/.cache/lax-ref +build/1bit npu-lax --model --kernels ~/.cache/lax/kernels --bench 256 [--transport classic] +build/1bit npu-lax --model --kernels ~/.cache/lax/kernels --snapshot-check [--dump ] +flock ctest --test-dir build -R npu_lax --output-on-failure +``` + +Without prompts on the command line, it reads one prompt per line from stdin, and the +conversation stays one device session. `--parity` runs the reference's `xres.bin` at +position t and scores the logits with `compare_decode.py`'s metric, then resets the decoder +(`Decoder::reset`: the DeltaNet state back to zero) and runs the positions again, which +must give bit-identical logits; `1bit serve` resets that way between conversations. +`--bench N` runs N tokens at positions 0..N-1, twice (new, then cached position configs), +and prints the host preparation, 40-layer and head times per token. + +The tests: + +| CTest | What it checks | Needs | +|---|---|---| +| `npu_lax_host` | The q8 -> q4_1 re-quantization and the signed-nibble transcode on synthetic chunks, every chunk permutation, the 4096-row position table, the position patches on a synthetic stream and the cfg words. Each is checked against hashes of the Python packer's output for the same input. | nothing (CI) | +| `npu_lax_pack_model` | All 83 buffers (ptab, final norm, lm head pool, 40 pools and 40 consts) against the SHA-256 of the Python packer's bytes, in `tests/golden/npu_lax/sha256.tsv` | the model | +| `npu_lax_host` (also) | The full-ELF position sites on a synthetic instruction ELF: which words and addends move, their values at a position, the fold kept until assembly clears it, and three malformed inputs refused | nothing (CI) | +| `npu_lax_patches` | The four position patches found in `lax_a/insts.bin` against the reference harness's table (`tests/golden/npu_lax/lax_a_patches.tsv`) | the kernels | +| `npu_lax_elf` | The seven full-ELF position sites of `lax_a/insts.elf` against the investigation's (`tests/golden/npu_lax/lax_a_elf_sites.tsv`); the PDI-less position configs equal the PDI ones in every other section; with `-DONEBIT_NPU_LAX_ELF_GOLDEN=`, all seven full ELFs byte for byte against the ones the reference harness ran | the kernels | +| `npu_lax_e2e` | `tests/npu_lax_cpp.sh`: on full ELFs, parity on three positions and again after a reset, then one chat turn that must answer with Paris | model, kernels, reference, the NPU | +| `npu_lax_e2e_classic` | the same on the classic path | model, kernels, reference, the NPU | +| `npu_lax_host` (also) | The chat follow-up bookkeeping (`npu/lax_turns.h`): the snapshot store's longest-proper-prefix match, LRU eviction and buffer reuse, the live-vs-snapshot plan, the snapshot points, and a three-turn chat replayed on a stand-in decoder whose state is a hash of the tokens fed | nothing (CI) | +| `npu_lax_snapshot` | `1bit npu-lax --snapshot-check`: a DeltaNet state snapshot restored after a different continuation gives logits bit-identical to the same prefix fed from scratch, and a `Session` follow-up from it generates the same tokens (see "Chat follow-ups") | model, kernels, the NPU | +| `npu_lax_serve` | `tests/npu_lax_serve.sh`: `1bit serve --device npu` under strace; the chat answers "The capital of France is Paris.", twice (the second from a snapshot), a chat follow-up reuses the first turn from a snapshot and answers Berlin, it streams with a usage chunk, and opens no `.xclbin` | model, kernels, the NPU | + +The reference hashes come from `model/lax_pack.py`'s `Model.build` at the pin. That +path produces the bytes `lax_chat.py` streams into the device, and they equal +`make_decode.py --requant`'s consts and ptab files. The q8 re-quantization has to round +the way NumPy does, step for step: + +- It is built without fused multiply-adds. +- A tied min or max goes to the later element. In a block whose scale is exactly zero, + every value is +0 or -0, and the min's sign is stored. + +**Results** on Strix Halo, with the same model and kernels as the Python numbers below. +The binary is `build/1bit` from branch `npu/lax-cpp` (the classic path; the full-ELF +results are under "Full ELFs"): + +- **Packer:** all 83 buffers are byte-identical to the Python packer. That covers the 40 + layers, the lm head and the ptab (`npu_lax_pack_model`, 17 s). +- **Parity** (`1bit npu-lax --parity`): corr 0.999998 / 0.999993 / 0.999998 and argmax + 846 / 198 / 3710 at positions 0 / 1 / 2. The top-5 equals the reference's. These are + the same numbers as the Python driver's. +- **Chat:** "The capital of France is Paris." + +Speed, measured back to back on the same prompt ("Write a detailed essay about the history +of Paris, from the Romans to today.", 28 prompt tokens, 128 generated). The generated text +is the same from both drivers: + +| Driver | Ready (process start to first prompt) | Prompt | Decode | +|---|---|---|---| +| `1bit npu-lax` (C++) | 3.8 s | 20.7 tok/s | 16.4 tok/s | +| `lax_chat.py` (Python) | 11.7 s | 18.7 tok/s | 16.2 tok/s | + +Decode is bound by the device: 40 layers take about 48 ms as one submit, and the lm head +streams 542 MB. The C++ driver gains its time at load. It fills 22.4 GB in 0.9 s over 16 +threads, after 2.8 s allocating the buffers. Both drivers were measured with +`model.q4nx` already in the page cache. + +## Kernel directory + +`scripts/build-lax.sh ` writes `/kernels`, which both transports read: + +``` +kernels/ + lax_l/ insts.elf main.pdi final.xclbin insts.bin final.prj/ linear-attention text + lax_a/ insts.elf main.pdi final.xclbin insts.bin final.prj/ full-attention text + ln/ insts.elf main.pdi final.xclbin insts.bin final RMSNorm + lm_head_q8/ insts.elf main.pdi final.xclbin insts.bin q8 lm head + manifest.json spec.json toolchain.json +``` + +Full ELFs read `insts.elf` and `main.pdi` (a kind without `main.pdi` falls back to +`final.prj/main.pdi`); the classic path reads `final.xclbin` and `insts.bin`. lax_l and +lax_a come from one design build, so lax_l's PDI configures the array for both. The export +of the pinned recipe copies only the classic files of ln and lm_head_q8, so the script +copies their `insts.elf` and PDI from the builds under `src/open_kernels/designs`. +`1bit serve` finds this directory through `--npu-kernels`, `/npu/lax` or +`$ONEBIT_NPU_LAX_KERNELS`. + +## Full ELFs + +Each kernel is assembled in memory at load time by `npu/full_elf` (`assemble_full_elf`, the +fast lane's assembler) from the build's `insts.elf` and PDI, and opened with +`xrt::hw_context(dev, xrt::elf)` / `add_config`. `strace` shows no `.xclbin` opened. +`npu/lax_elf.h` builds them, and `npu/lax_elf_kernels.cpp` loads them: + +| Kernel | From | Context | +|---|---|---| +| `laxinit` | lax_l's PDI, only `load_pdi`, one unreferenced buffer argument | the layer context is created from it | +| `lxf` | lax_l's text, no `load_pdi` | a config of the layer context | +| `axf` | lax_a's text written for position pos, no `load_pdi`, no PDI | a config of the layer context, added the first time the decode reaches pos | +| `ln`, `lm` | `load_pdi` + text, their own PDIs | one stand-alone context each | + +Buffers go from argument 0: `lxf`/`axf` `pool xres consts kv act ptab state cfg` (lxf gets +the unused 8 MB kv, axf the unused 2.3 MB state), `ln` `xres zero normw xresf hn`, `lm` +`lmpool hn logits`. `cfg` words 0..1 stay the pool's address + 0x80000000. + +**Why the ELF path used to fault.** mlir-aie folds 0x80000000 (`kDDRAIEAddrOffset`) into +the DDR_PATCH offset of buffer arguments >= 5, for the classic firmware, which translates +only the first five. aiebu-asm moves that offset into the `.rela.dyn` addend, and XRT's +ELF patcher adds `bo.address() + 0x80000000` for every argument. So `ptab`, `state` and +`cfg` were translated twice. `unfold()` (`npu/lax_elf.h`) clears the fold, as +`assemble_full_elf`'s bit-31 mask already did, and refreshes the UID note to the md5 of +the control text XRT runs. This was found on OpenFlowLM-Next branch `1bit/lax-elf` +(65ad44d, `35b-whole-layer-runlist-status.md`). The pin does not need that branch: the +engine unfolds the default build. + +**Every token's runlist starts with the `laxinit` run.** ln and lm_head run in their own +contexts, and after them the layer context's array configuration is gone: with init once, +token 1 timed out. Heading each runlist with it costs nothing measurable. + +**Positions.** On the ELF path the full-attention position is not a word the host can +rewrite: XRT patches the BD from the relocation addend. `elf_position_sites` +(`npu/lax_stream.h`) derives the seven places from `insts.elf`. It finds the same four +DDR_PATCH patches as the classic `position_patches`. For the window fill, the row drain and +the RoPE record, it takes the DDR_PATCH word and the addend of the one relocation on the +BD that patch rewrites (its symbol must name the same argument, and its addend must equal +the word). The window length is the fourth site. For the pinned build these are ctrl words +3222 (window length, `max(pos, 1) * 512`), 2762 + relocation 68 (drain, `pos * 2048`), +2984 + relocation 74 (record, `pos * 1024`) and 3240 + relocation 78 (window offset, 0). +Each position is its own config: + +- Adding one takes 0.2-0.5 ms. +- The decoder prepares the next position's runlist (its config, the ten full-attention + runs, the runlist) while the device runs the current one, so none of it is on the + token's critical path. +- Position configs carry no PDI. They have no `load_pdi`, so XRT never uploads one, but it + keeps every added ELF. With lax_l's 257 KiB PDI in each, 128 positions raised the + maximum RSS by 34 MB over the classic path (about 1 GB at 4096 positions). Without it, + 512 positions stayed within 13 MB of the classic path (21.98 vs 21.96-21.97 GB). + +**Results** on Strix Halo, 2026-09-24: + +- Model: `model.q4nx` sha256 `688f1e153d10…3cf8de`. +- Kernels: `scripts/build-lax.sh` at the pin, with `insts.bin` md5 `136cc0c9…` / `4f3c749d…` + and lax_a `insts.elf` md5 `426b8dd1…`. +- Binary: `build/1bit` from branch `npu/lax-elf`. + +The results: + +- **Assembly** (`npu_lax_elf` with the harness's ELFs): lax_init, lxf, axf_p0/1/2, ln + and lm are byte-identical to the ones `harness/full_elf.py lax` assembled and ran. +- **Parity** (`1bit npu-lax --parity`, full ELFs): corr 0.999998 / 0.999993 / 0.999998 and + argmax 846 / 198 / 3710 at positions 0 / 1 / 2, the same top-5 as the reference, and + bit-identical again after a reset. The logits are byte-identical to the classic path's + (`--dump`, `cmp`). +- **Serve** (`tests/npu_lax_serve.sh`): "The capital of France is Paris.", no `.xclbin` + opened. + +Speed, A/B in alternating runs of the same binary, host load average 1-5: + +| | Full ELF | Classic | +|---|---|---| +| `--bench 256`, per token (40 layers + norm + lm head) | 59.7-60.8 ms | 61.2-62.1 ms | +| of which the 40-layer runlist | 49.2-50.0 ms | 50.6-51.3 ms | +| host preparation before the submit | 0.001-0.003 ms (plus 0.8-1.7 ms inside the runlist's time, preparing the next position) | 0.004 ms | +| Chat, "Write a detailed essay about the history of Paris, from the Romans to today.", 128 tokens: decode | **16.3 / 16.5 tok/s** | 16.0 / 16.1 tok/s | +| the same: prompt (28 tokens) | 20.8 / 20.9 tok/s | 20.4 / 20.3 tok/s | +| Ready (process start to first prompt) | 3.6-3.8 s | 3.6 s | + +Both transports generate the same text. Before the lookahead, building the next +runlist after the submit cost 1.1-1.4 ms per token, as much as the ELF path saves. + +## Chat follow-ups (state snapshots) + +A chat client sends the whole history with every request. Qwen3.6's chat template +re-renders each earlier answer without the think block it was generated after (with +thinking off, without the empty `\n\n\n\n`), so a follow-up never extends +the tokens on the device, and before this it started over at position 0. + +A position's cache has two parts: + +- **The KV rows** of the 10 full-attention layers (2048 B per position, `[K | V]`). A token + at position `pos` streams rows `0..pos-1` (the window length is `max(pos, 1)` rows; at 0 + the one row is masked) and writes row `pos`. So after a jump back to any position, the + stale rows beyond it are overwritten before anything reads them. +- **The DeltaNet state** of the 30 linear layers (recurrent S plus conv state, one + 2,342,912 B buffer each). Every token is folded into it, so it cannot rewind. + +`lax::Session` (`npu/lax.h`, with the bookkeeping in `npu/lax_turns.{h,cpp}`) keeps +**snapshots** of the DeltaNet state: `Decoder::save_state` copies the 30 buffers into host +memory (70,287,360 B), and `load_state` copies them back. It takes a snapshot while a prompt +is fed, at two prefixes: + +- **Before the prompt's last token.** This covers a retry of the same prompt, or a prompt + that extends it but not its answer. +- **Before the prompt's last `<|im_start|>`**, the assistant header. A follow-up repeats + every message before it exactly. + +A request starts from whichever leaves the fewest tokens to feed: + +- the live cache, when the prompt extends the tokens on the device; +- the longest snapshot whose tokens are a proper prefix of the prompt: load it, and continue + at the position of its length; +- otherwise, a reset and position 0. + +The last prompt token always runs, for its logits. Snapshots are kept least recently used +first out. `1bit serve --npu-snapshots N` sets how many (default 4, 281 MB; 0 turns them +off). `usage.prompt_tokens_details.cached_tokens` reports the reused tokens, and +`timings.cache` reports where they came from (`live` or `snapshot`). + +**Correctness** (`npu_lax_snapshot`, `1bit npu-lax --snapshot-check`): + +1. Feed P, then X from scratch, and keep X's logits. +2. Reset, feed P and save the state. Feed a different continuation Y, which rewrites the + state, the layers' scratch and the KV rows at X's positions. +3. Load the state and feed X at the same positions. + +X's logits must be the same bytes. Then the same runs through `Session`: P + Y, then +P + X from the snapshot, must generate what P + X generates from scratch. + +Results on Strix Halo, 2026-09-24, with `model.q4nx` sha256 `688f1e153d10…3cf8de`, the +kernels under "Full ELFs", and `build/1bit` from branch `npu/lax-turns`: + +- **Logits:** bit-identical at all 28 of X's positions. P was 17 tokens, X 28 and Y 23. + `--dump` wrote `scratch.bin` and `restored.bin` (27,811,840 B each), and `cmp` finds them + identical. +- **Session:** the follow-up reused 17 tokens from the snapshot and generated the same 8 + tokens ("The capital of Germany is Berlin."). +- **Timing:** `save_state` took 4.7 ms and `load_state` 5.1 ms. + +Only the DeltaNet buffers need restoring. The layers' `act` scratch is left as Y wrote it, +and the logits are still bit-identical. + +**A three-turn chat** (`tests/npu_lax_turns.py` against `1bit serve --device npu`): OpenAI +messages with the whole history, each answer sent back as the client got it, thinking off, +`max_tokens` 64, streamed with `stream_options.include_usage`. The first turn is a +197-token prompt (trip notes and a question). Time to first token is measured by the +client. Baseline `before`: `build/1bit` at `7c75672` (PR #52's head), which had no usage chunk +in streams, so its times are the server's `timings.prompt_ms` from `--no-stream`. For the new +binary those equal its streamed time to first token within 0.01 s. `--npu-snapshots 0` on +the new binary matched the baseline to within 0.1 s. + +| Turn | Prompt tokens | Fed again before | Fed again after | Time to first token before | after | +|---|---|---|---|---|---| +| 1 | 197 | 197 | 197 | 9.50 s | 9.46 s | +| 2 | 267 | 267 | **77** (190 from a snapshot) | 12.83 s | **3.78 s** | +| 3 | 322 | 322 | **62** (260 from a snapshot) | 15.48 s | **3.09 s** | + +- **Output:** the answers are the same text before and after. +- **Per request:** two snapshots cost 9-14 ms, and a restore 5 ms. +- **Memory:** 4 snapshots take 281 MB of host memory (70.3 MB each). The serve process's + RSS was 22.26-22.28 GB with snapshots and 22.01 GB with `--npu-snapshots 0`. +- **The tokens still fed:** each follow-up feeds the previous answer and the new question. + That answer was generated after a think header the history leaves out, so its state + cannot be reused. + +## Results + +All results are on Strix Halo, using `model.q4nx` sha256 `688f1e153d10…3cf8de` (the +`Qwen3.6-35B-A3B-NPU2` Q4NX container, 23,235,412,536 B). The pin is `2490fa6e6ddb`, +built by `scripts/build-lax.sh`. It produces `insts.bin` md5 `136cc0c9…` (lax_l) and +`4f3c749d…` (lax_a). + +**Parity** (`tests/npu_lax_parity.sh`), three positions from `<|im_start|>`, each +position's greedy token fed back: + +| position | logits corr | argmax (ours = reference) | top-5 | +|---|---|---|---| +| 0 | 0.999998 | 846 | identical | +| 1 | 0.999993 | 198 | identical | +| 2 | 0.999998 | 3710 | identical | + +The chat turn answers "The capital of France is Paris." + +**Speed** (`lax_chat.py`, 128 generated tokens, host load average 8-10): + +- 16.1 tok/s end to end. That covers the 40 layers as one submit, the norm, the lm + head, and the argmax fed back to the host. +- Prompt tokens run at about 19 tok/s. +- Load, from process start to ready (packing 21.8 GB into the buffers): 12.1 s. + +### How it got there + +| Step (commit on `1bit/lax-35b`) | 40 layers | tok/s | +|---|---|---| +| First correct decode (469c43b) | 75.8 ms | 11.1 | +| Full-attention layers skip the dummy DeltaNet and pad bands (3632f7d) | 67.5 ms | 12.3 | +| DeltaNet fills and drains issued head-major across the 8 cores (0866c3f) | 53.5 ms | 14.7 | +| Router weights on two streams, vector top-8, MoE header overlapped (2490fa6) | 47.7-48.1 ms | 16.1 | + +### What made it correct + +Before 469c43b, the design had only ever completed on unfilled buffers. On real +weights it returned NaN. Three faults were behind that: + +1. **Two `acquire(1)` calls return the same element.** The emitter core took the + router output and its config as two `acquire(1)` calls on the same object fifo. + Acquire counts are cumulative, so the second call returned the first element again. + The pool base address was therefore read from the router's probabilities. Fix: + `acquire(2)`. +2. **The host and the emitters raced on one queue.** The host queued the shared + expert's fills on the same MM2S queue the emitters push the routed experts to. The + host did not wait for the last routed push, and the fills are the same size, so the + core silently used the wrong weights. Fix: the emitter pushes the shared expert too + (`ONDV_EMIT_SHARED`). +3. **A lock was released but never acquired.** It counted up by one per layer, and + AIE2 locks are 6-bit, so the device hung after 8 layers per context. This was the + "8 layers per context" cap. Fix: the emitter acquires the lock back + (`ONDV_PKTDONE_ACQ`). + +### Where the time goes + +A per-stage profile of one layer at 469c43b, taken with on-device timestamps, found +two things: + +- The weight streams already run at 45-53 GB/s, close to a measured DDR read peak of + about 57 GB/s. +- The losses were a staggered DeltaNet schedule, dead work in the full-attention + layers, and the router stage. + +The last three steps in the table remove those losses. The DDR bandwidth limit for a +token is about 22 tok/s (estimate). + +### Compared with the GPU + +On the same box, the Q8_0 GGUF of the model does: + +| Backend | Decode (tg128) | Prompt (pp512) | Binary | +|---|---|---|---| +| Vulkan | 53.4 tok/s | 1269 tok/s | `llama-bench`, llama.cpp 7fe450e | +| HRX0 | 40.9 tok/s | fails: memory fault on batches of 2 or more | `llama-bench`, llama.cpp f1a0aca / hrx-system 51b1739 | + +Both used `-ngl 99 -fa 1 -r 3`. The GPU is the faster path for this model. What the +NPU adds is a second, independent decode stream. At 469c43b the NPU decoded 11.0 tok/s +next to a running Vulkan benchmark without slowing it (as noise measures). It also +draws about 75-80 W less package power than GPU decode. + +## What is left + +- **Prompt processing.** Prompt tokens go through one position at a time. There is no + batched prefill on the NPU. +- **Memory.** A session pins about 22 GB of host memory in XRT buffers. +- **The previous answer is fed again.** A chat follow-up restores the snapshot taken before + the previous prompt's assistant header ("Chat follow-ups"). It then feeds the answer + again, re-rendered without its think block, plus the new question, at about 21 tok/s. + Reusing the answer's own state would need a snapshot at a position the re-rendered + history also reaches. No such position exists: the answer was generated after the + think header, and the re-rendered history leaves that header out. +- **Snapshots live in host memory only.** They do not survive a restart. With the default + 4, a client that switches between more than two conversations evicts the others' + snapshots. diff --git a/docs/serve.md b/docs/serve.md index c23f8f04..7691980c 100644 --- a/docs/serve.md +++ b/docs/serve.md @@ -29,6 +29,7 @@ OpenAI client. [--mtp HEAD.gguf] [--mtp-max N] [--mtp-p-min P] [--parallel N] [--adaptive] [--adaptive-at N] [--embed MODEL.gguf] [--rerank MODEL.gguf] + [--npu-kernels DIR] [--npu-transport elf|classic] [--npu-snapshots N] ``` One model per process: @@ -47,6 +48,7 @@ One model per process: | Model | `--device` | Runs on | |---|---|---| | NPU model directory (`model.q4nx` + `npu/`, docs/npu.md) | `auto`, `npu` | the NPU fast lane, in process | +| Qwen3.6-35B-A3B Q4NX directory (`model_type` `qwen3_5_moe`) with the lax kernels | `auto`, `npu` | the NPU lax decode on full ELFs, in process (docs/npu-lax.md) | | `.gguf` | `auto`, `vulkan` | the upstream llama.cpp build's llama-server on `Vulkan0` (docs/vulkan.md); without `ONEBIT_VULKAN`, the HRX build's | | `.gguf` | `hrx` | the HRX build's llama-server on `HRX0` (docs/hrx.md) | | `.gguf` | `vulkan --prefill-device hrx` | the HRX build's llama-server decoding on `Vulkan0`, long prompt prefixes prefilled on `HRX0` over one shared KV cache (docs/hrx.md, "Prefill on HRX, decode on Vulkan") | @@ -55,9 +57,34 @@ One model per process: | `.gguf` | `zinc` | this build's ZINC (Vulkan, ROCm or CUDA, whichever it was built for; docs/zinc.md) | | Hugging Face id | `mlx` | lemon-mlx-engine's server, on Apple Silicon (docs/apple.md) | +The lax decode needs the kernels `scripts/build-lax.sh` builds (docs/npu-lax.md, "Kernel +directory"). `serve` looks for them in `--npu-kernels DIR`, else `/npu/lax`, else +`$ONEBIT_NPU_LAX_KERNELS`; a Qwen3.6-35B-A3B directory without them is not an NPU model +directory. `--npu-transport classic` runs the same kernels through the xclbin path, for A/B. +The decode keeps one conversation on the device, and snapshots of its DeltaNet state in +host memory (docs/npu-lax.md, "Chat follow-ups"). A request starts from whichever leaves +the fewest tokens to feed: + +- the device's tokens, when its prompt extends them; +- a snapshot at a prefix of its prompt. Snapshots are taken before each prompt's last + token and before its last `<|im_start|>`, so a chat client's follow-up re-feeds only the + previous answer and the new question; +- otherwise, position 0. + +`--npu-snapshots N` keeps up to N snapshots, 70 MB each (default 4; 0 reuses only exact +extensions). `usage.prompt_tokens_details.cached_tokens` counts the reused tokens, and +`timings.cache` says where they came from (`live` or `snapshot`). On this route a streamed +request with `stream_options.include_usage` gets a last chunk with `usage` and `timings`, +as OpenAI's API sends it. + +```sh +1bit serve -m ~/models/Qwen3.6-35B-A3B-NPU2 --device npu --npu-kernels ~/.cache/lax/kernels +``` + `chat_template_kwargs.enable_thinking: false` works on every device. The GPU backends apply the model's own chat template. The NPU route emits what Qwen3's -template does: an empty think block after the assistant prefix. +template does: an empty think block after the assistant prefix. For Qwen3.6-35B-A3B it +also opens `\n` when thinking is on, as that model's template does. For a `.gguf` the engine starts that server as a private child on a loopback port and forwards the OpenAI routes to it, streaming included. Replies carry @@ -220,6 +247,39 @@ streaming: | `hrx` | PASS (32 SSE chunks), with no environment set up | | `zinc` | PASS (6 SSE chunks) | +The 35B on the NPU: `tests/npu_lax_serve.sh` (ctest `npu_lax_serve`, with +`-DONEBIT_NPU_LAX_MODEL` and `-DONEBIT_NPU_LAX_KERNELS`) serves Qwen3.6-35B-A3B under +strace. On Strix Halo, 2026-09-24, `build/1bit` from branch `npu/lax-elf`, `model.q4nx` +sha256 `688f1e153d10…3cf8de`: + +- Ready in 4.1-4.3 s with `model.q4nx` in the page cache (7.5 s once, when it was not). +- The chat answers "The capital of France is Paris." (24 prompt tokens in 1.13 s, 7 + generated in 0.42 s). The same request again answers the same after the decode starts + over. +- A raw prompt that extends that turn reuses all 31 of its tokens + (`usage.prompt_tokens_details.cached_tokens`) and answers "The capital of Germany is + Berlin." +- The same follow-up sent as chat messages reuses none of its 45 tokens, because the + template renders the earlier answer without its think block. +- 15 SSE chunks stream, and no `.xclbin` is opened (the four `insts.elf` are). + +PASS. + +With state snapshots, on the same day and model, `build/1bit` from branch `npu/lax-turns` +gave the results in the list below. docs/npu-lax.md ("Chat follow-ups") has the three-turn chat: time to first +token went from 12.83 s to 3.78 s (turn 2) and from 15.48 s to 3.09 s (turn 3). + +- The same request again restores the snapshot before its last token: 23 of 24 tokens are + cached, the prompt takes 66 ms instead of 1172 ms, and the answer is the same. +- The chat follow-up reuses 17 of its 45 tokens from the snapshot before the first prompt's + assistant header (`timings.cache` `snapshot`). It answers "The capital of Germany is + Berlin." +- The raw prompt that extends a turn still reuses all 31 of that turn's tokens from the live + cache. +- A streamed request with `stream_options.include_usage` ends with the usage chunk. + +PASS. + ctest runs these as `serve_e2e_` when configured with `-DONEBIT_SERVE_TEST_GGUF=` (and `serve_e2e_mlx` on macOS with `-DONEBIT_MLX_SERVER`). `smoke_serve` runs everywhere, CI included: diff --git a/npu/full_elf.cpp b/npu/full_elf.cpp index dfe30a30..82b1adb5 100644 --- a/npu/full_elf.cpp +++ b/npu/full_elf.cpp @@ -81,6 +81,12 @@ size_t uid_offset(const Bytes& e) { } // namespace +ElfSection elf_section(const Bytes& elf, const std::string& name) { + const Section s = section(sections(elf), name.c_str()); + if (uint64_t(s.offset) + s.size > elf.size()) fail(name + " runs past the end of the file"); + return {s.offset, s.size}; +} + Bytes read_file(const std::string& path) { std::ifstream f(path, std::ios::binary); if (!f) throw std::runtime_error("cannot read " + path); @@ -115,10 +121,14 @@ Bytes derive_context(const Bytes& ctx1, const ContextMap& map, int n) { const uint32_t blocks = uint32_t(n + ContextMap::kBlock - 1) / ContextMap::kBlock; for (const auto& w : map.words) wr32(e, w.offset, w.step ? w.base + uint32_t(n - 1) * w.step : w.base * blocks); + refresh_uid(e); + return e; +} + +void refresh_uid(Bytes& e) { const Section ct = section(sections(e), ".ctrltext"); const auto digest = md5({e.data() + ct.offset, ct.size}); std::memcpy(e.data() + uid_offset(e), digest.data(), 16); - return e; } // ---- Full-ELF assembly ---- diff --git a/npu/full_elf.h b/npu/full_elf.h index cfd3126c..89630d6a 100644 --- a/npu/full_elf.h +++ b/npu/full_elf.h @@ -34,6 +34,13 @@ using Bytes = std::vector; Bytes read_file(const std::string& path); +// A section of an ELF32 little-endian file, by name: its byte offset and size in the +// file. Throws if the file is not an ELF32 LE file or has no such section. +struct ElfSection { + uint32_t offset, size; +}; +ElfSection elf_section(const Bytes& elf, const std::string& name); + // How the layer kernel's control code depends on the context length N. // Every differing word follows one of two rules, found by comparing the // captured contexts 1, 2 and 17: @@ -58,6 +65,10 @@ struct ContextMap { // captured layer_ctx.elf. Bytes derive_context(const Bytes& ctx1, const ContextMap& map, int n); +// Rewrite an instruction ELF's UID note to the md5 of its .ctrltext, as aiebu-asm writes +// it, after the control code was edited. +void refresh_uid(Bytes& instruction_elf); + enum class PdiMode { kLoad, // control code starts with load_pdi (a stand-alone kernel) kNone, // no load_pdi: runs on an array an init config already configured diff --git a/npu/lax.cpp b/npu/lax.cpp new file mode 100644 index 00000000..1c229b19 --- /dev/null +++ b/npu/lax.cpp @@ -0,0 +1,525 @@ +// 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 program this drives is the one open_kernels' model/lax_decode_cfg.py (setup_packed, +// position) writes for its XRT harness (third_party/OpenFlowLM-Next, MIT): the same +// buffers, argument order and per-token steps, re-implemented against XRT directly. +#include "lax.h" + +#include "lax_kernels.h" +#include "lax_stream.h" + +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace onebit::npu::lax { + +namespace { + +using clk = std::chrono::steady_clock; +double ms_since(clk::time_point t) { return std::chrono::duration(clk::now() - t).count(); } + +constexpr size_t kBoAlign = 1u << 20; // XDNA buffers in whole MiB, as the harness sizes them +size_t padup(size_t n) { return (n + kBoAlign - 1) / kBoAlign * kBoAlign; } +constexpr auto kTimeout = std::chrono::milliseconds(60000); + +std::vector slurp(const std::string& path) { + std::ifstream f(path, std::ios::binary | std::ios::ate); + if (!f) throw std::runtime_error("cannot read " + path); + std::vector v(size_t(f.tellg())); + f.seekg(0); + if (!v.empty() && !f.read(reinterpret_cast(v.data()), std::streamsize(v.size()))) + throw std::runtime_error("short read " + path); + return v; +} + +// ---- the classic kernel set -------------------------------------------------------------- + +class Classic final : public KernelSet { + struct Stream { + xrt::kernel kernel; + xrt::bo instr; + size_t nwords = 0; + uint32_t* words() { return instr.map(); } + }; + +public: + Classic(xrt::device& dev, const std::string& dir) : dev_(dev), dir_(dir) { + // Both control texts come from one design build, so both run on one context and one + // runlist: lax_l's xclbin serves lax_a's stream too, as the reference driver does + // (lax_a's own xclbin differs only in its UUID and metadata). + layer_ctx_ = context(dir + "/lax_l/final.xclbin"); + ln_ctx_ = context(dir + "/ln/final.xclbin"); + lm_ctx_ = context(dir + "/lm_head_q8/final.xclbin"); + lin_ = stream(layer_ctx_, dir + "/lax_l/insts.bin"); + full_ = stream(layer_ctx_, dir + "/lax_a/insts.bin"); + ln_ = stream(ln_ctx_, dir + "/ln/insts.bin"); + lm_ = stream(lm_ctx_, dir + "/lm_head_q8/insts.bin"); + patches_ = position_patches({full_.words(), full_.nwords}, kKvRow); + } + + xrt::hw_context& layer_context() override { return layer_ctx_; } + + xrt::run make_run(Stage s) override { + Stream& st = s == Stage::LayerLinear ? lin_ : s == Stage::LayerFull ? full_ : s == Stage::Norm ? ln_ : lm_; + xrt::run r(st.kernel); + r.set_arg(0, 3); // opcode: run the instruction stream + r.set_arg(1, st.instr); + r.set_arg(2, int(st.nwords)); + return r; + } + + int first_buffer_arg() const override { return 3; } + + void set_position(size_t pos) override { + apply_position({full_.words(), full_.nwords}, patches_, pos, kKvRow, kPtabRow); + full_.instr.sync(XCL_BO_SYNC_BO_TO_DEVICE); + } + + bool position_in_run() const override { return false; } + + std::optional head_run() override { return std::nullopt; } + + std::string describe() const override { return "classic xclbin + insts.bin from " + dir_; } + +private: + xrt::hw_context context(const std::string& path) { + const xrt::xclbin x(path); + return xrt::hw_context(dev_, dev_.register_xclbin(x)); + } + + Stream stream(xrt::hw_context& ctx, const std::string& path) { + const auto bytes = slurp(path); + if (bytes.empty() || bytes.size() % 4) throw std::runtime_error(path + ": not whole instruction words"); + Stream s{xrt::kernel(ctx, "MLIR_AIE"), {}, bytes.size() / 4}; + s.instr = xrt::bo(dev_, bytes.size(), xrt::bo::flags::cacheable, s.kernel.group_id(1)); + std::memcpy(s.instr.map(), bytes.data(), bytes.size()); + s.instr.sync(XCL_BO_SYNC_BO_TO_DEVICE); + return s; + } + + xrt::device& dev_; + std::string dir_; + xrt::hw_context layer_ctx_, ln_ctx_, lm_ctx_; + Stream lin_, full_, ln_, lm_; + std::vector patches_; +}; + +} // namespace + +std::unique_ptr classic_kernels(xrt::device& dev, const std::string& dir) { + return std::make_unique(dev, dir); +} + +// ---- the decoder ----------------------------------------------------------------------- + +struct Decoder::Impl { + using Bo = std::unique_ptr; + const Model& model; + const Config cfg; + xrt::device dev{0u}; + std::unique_ptr ks; + Bo xres, zero, normw, xresf, hn, logits, lmpool, ptab, dkv, dstate; + std::vector pool, consts, act, lcfg, cache; // cache: DeltaNet state or KV, per layer kind + // A token's runlist, for one position. With full ELFs, the next position's slot is + // prepared while the device runs the current one (the kernel set's position_in_run); + // classic uses slot 0 only and patches its stream in place. + struct Slot { + std::vector runs; // one per layer, in the runlist's order + std::unique_ptr rl; + size_t pos = SIZE_MAX; + }; + std::vector linear_runs; // position-free, shared by both slots (empty for full layers) + Slot slots[2]; + std::optional ln, lm; + const uint8_t* embed = nullptr; + LoadStats stats; + RunTimes times; + + Bo bo(size_t bytes, bool zeroed = true) { + auto b = std::make_unique(dev, padup(bytes)); + if (zeroed) { + std::memset(b->map(), 0, padup(bytes)); + b->sync(XCL_BO_SYNC_BO_TO_DEVICE); + } + return b; + } + + Impl(const Model& m, const Config& c, const std::string& dir, Transport t, int threads) : model(m), cfg(c) { + const auto t0 = clk::now(); + ks = make_kernels(dev, dir, t); + const Tensor& et = m.tensor("model.embed_tokens.weight"); + if (et.dtype != "BF16" || et.bytes != size_t(c.vocab) * kHidden * 2) throw std::runtime_error("unexpected embedding table"); + embed = et.data; + + // Buffers. The packed ones are written whole by their packer (no zero pass). + const auto t1 = clk::now(); + xres = bo(kHidden * 4); + zero = bo(kHidden * 4); + normw = bo(kHidden * 2); + xresf = bo(kHidden * 4); + hn = bo(kHidden * 2); + logits = bo(kLogitsBytes); + dkv = bo(kKvBytes); + dstate = bo(kStateBytes); + lmpool = bo(kLmPoolBytes, false); + ptab = bo(kPtabBytes, false); + for (int L = 0; L < c.layers; ++L) { + const Kind k = c.kinds[size_t(L)]; + pool.push_back(bo(kPoolBytes, false)); + consts.push_back(bo(kConstsBytes, false)); + act.push_back(bo(kActBytes)); + lcfg.push_back(bo(kCfgBytes)); + cache.push_back(bo(k == Kind::Linear ? kStateBytes : kKvBytes)); + } + stats.buffers_ms = ms_since(t1); + + // Pack in parallel, each job straight into its buffer's host mapping. + const auto t2 = clk::now(); + std::vector> jobs; + auto sync = [](xrt::ext::bo& b) { b.sync(XCL_BO_SYNC_BO_TO_DEVICE); }; + jobs.emplace_back([&] { + pack_lmhead(model, static_cast(lmpool->map())); + std::memset(static_cast(lmpool->map()) + kLmPoolBytes, 0, padup(kLmPoolBytes) - kLmPoolBytes); + sync(*lmpool); + }); + jobs.emplace_back([&] { + pack_ptab(cfg, kMaxContext, static_cast(ptab->map())); + sync(*ptab); + pack_norm(model, static_cast(normw->map())); + sync(*normw); + }); + for (int L = 0; L < c.layers; ++L) { + jobs.emplace_back([&, L] { + pack_pool(model, cfg, L, static_cast(pool[size_t(L)]->map())); + sync(*pool[size_t(L)]); + }); + jobs.emplace_back([&, L] { + auto* d = static_cast(consts[size_t(L)]->map()); + const size_t n = consts_bytes(cfg.kinds[size_t(L)]); + pack_consts(model, cfg, L, d); + std::memset(d + n, 0, padup(kConstsBytes) - n); + sync(*consts[size_t(L)]); + }); + } + if (threads <= 0) threads = int(std::clamp(std::thread::hardware_concurrency(), 1u, 16u)); + std::atomic next{0}; + std::exception_ptr err; + std::mutex err_mu; + std::vector pool_threads; + for (int t = 0; t < threads; ++t) + pool_threads.emplace_back([&] { + for (size_t i; (i = next++) < jobs.size();) { + try { + jobs[i](); + } catch (...) { + std::lock_guard lk(err_mu); + if (!err) err = std::current_exception(); + next = jobs.size(); + } + } + }); + for (auto& t : pool_threads) t.join(); + if (err) std::rethrow_exception(err); + stats.pack_ms = ms_since(t2); + stats.packed_bytes = kLmPoolBytes + kPtabBytes + size_t(c.layers) * kPoolBytes; + for (int L = 0; L < c.layers; ++L) stats.packed_bytes += consts_bytes(c.kinds[size_t(L)]); + + // Layer configs: the pool's device address and the weight-stream queues. + for (int L = 0; L < c.layers; ++L) { + const auto w = cfg_words(pool[size_t(L)]->address() + 0x80000000ull); + std::memcpy(lcfg[size_t(L)]->map(), w.data(), sizeof w); + lcfg[size_t(L)]->sync(XCL_BO_SYNC_BO_TO_DEVICE); + } + + // The token's runlist: 40 runs on one context, submitted once per token. + for (int L = 0; L < c.layers; ++L) + linear_runs.push_back(c.kinds[size_t(L)] == Kind::Linear ? layer_run(size_t(L)) : xrt::run()); + prepare(slots[0], 0); + const int a = ks->first_buffer_arg(); + ln = ks->make_run(Stage::Norm); + for (int i = 0; auto* b : {&xres, &zero, &normw, &xresf, &hn}) ln->set_arg(a + i++, **b); + lm = ks->make_run(Stage::Head); + for (int i = 0; auto* b : {&lmpool, &hn, &logits}) lm->set_arg(a + i++, **b); + stats.total_ms = ms_since(t0); + } + + // A new sequence: the DeltaNet state and the layers' scratch back to zero, as at load. + // The KV cache needs nothing: a position reads only the rows written before it. + void reset() { + for (int L = 0; L < cfg.layers; ++L) { + for (auto* b : {&act[size_t(L)], cfg.kinds[size_t(L)] == Kind::Linear ? &cache[size_t(L)] : nullptr}) { + if (!b) continue; + std::memset((*b)->map(), 0, (*b)->size()); + (*b)->sync(XCL_BO_SYNC_BO_TO_DEVICE); + } + } + } + + // The linear layers' state buffers, in layer order: what save_state/load_state copy. + size_t state_bytes() const { + return size_t(std::ranges::count(cfg.kinds, Kind::Linear)) * kStateBytes; + } + + void save_state(uint8_t* out) { + for (int L = 0; L < cfg.layers; ++L) { + if (cfg.kinds[size_t(L)] != Kind::Linear) continue; + auto& b = *cache[size_t(L)]; + b.sync(XCL_BO_SYNC_BO_FROM_DEVICE, kStateBytes, 0); + std::memcpy(out, b.map(), kStateBytes); + out += kStateBytes; + } + } + + void load_state(const uint8_t* in) { + for (int L = 0; L < cfg.layers; ++L) { + if (cfg.kinds[size_t(L)] != Kind::Linear) continue; + auto& b = *cache[size_t(L)]; + std::memcpy(b.map(), in, kStateBytes); + b.sync(XCL_BO_SYNC_BO_TO_DEVICE, kStateBytes, 0); + in += kStateBytes; + } + } + + void write_xres() { xres->sync(XCL_BO_SYNC_BO_TO_DEVICE, kHidden * 4, 0); } + + void wait(xrt::run& r, const char* what) { + r.start(); + if (r.wait(kTimeout) != ERT_CMD_STATE_COMPLETED) throw std::runtime_error(std::string(what) + " run did not complete"); + } + + // Layer L's run: pool xres consts kv|dkv act ptab state|dstate cfg (a linear layer gets + // the dummy kv, a full one the dummy state). + xrt::run layer_run(size_t l) { + const bool lin = cfg.kinds[l] == Kind::Linear; + xrt::run r = ks->make_run(lin ? Stage::LayerLinear : Stage::LayerFull); + const int a = ks->first_buffer_arg(); + r.set_arg(a + 0, *pool[l]); + r.set_arg(a + 1, *xres); + r.set_arg(a + 2, *consts[l]); + r.set_arg(a + 3, lin ? *dkv : *cache[l]); + r.set_arg(a + 4, *act[l]); + r.set_arg(a + 5, *ptab); + r.set_arg(a + 6, lin ? *cache[l] : *dstate); + r.set_arg(a + 7, *lcfg[l]); + return r; + } + + // Slot s runs position pos: the full-attention runs made for it (full ELFs), or the + // shared stream patched (classic); the head run first when the kernel set has one. + void prepare(Slot& s, size_t pos) { + ks->set_position(pos); + if (s.rl && !ks->position_in_run()) { + s.pos = pos; + return; + } + s.pos = SIZE_MAX; + s.runs.clear(); + for (size_t l = 0; l < linear_runs.size(); ++l) + s.runs.push_back(cfg.kinds[l] == Kind::Linear ? linear_runs[l] : layer_run(l)); + if (!s.rl) s.rl = std::make_unique(ks->layer_context()); + s.rl->reset(); + if (auto h = ks->head_run()) s.rl->add(*h); + for (auto& r : s.runs) s.rl->add(r); + s.pos = pos; + } + + void run(size_t pos, bool head) { + if (pos >= size_t(kMaxContext)) throw std::runtime_error("position " + std::to_string(pos) + " past the context"); + const auto t0 = clk::now(); + times = {}; + const bool ahead = ks->position_in_run(); + Slot* s = &slots[0]; + if (ahead && slots[1].pos == pos) s = &slots[1]; + else if (!ahead || slots[0].pos != pos) prepare(*s, pos); + times.prep_ms = ms_since(t0); + const auto t1 = clk::now(); + s->rl->execute(); + if (ahead && pos + 1 < size_t(kMaxContext)) { + // The next position, while the device runs this one; a failure here only means + // the next run() prepares it itself. + Slot& o = s == &slots[0] ? slots[1] : slots[0]; + const auto t = clk::now(); + try { + prepare(o, pos + 1); + } catch (const std::exception&) { + o.pos = SIZE_MAX; + } + times.ahead_ms = ms_since(t); + } + if (s->rl->wait(kTimeout) == std::cv_status::timeout) throw std::runtime_error("layer runlist timed out"); + times.layers_ms = ms_since(t1); + if (head) { + const auto t2 = clk::now(); + wait(*ln, "norm"); + wait(*lm, "lm head"); + times.head_ms = ms_since(t2); + } + } + + const float* read_logits() { + logits->sync(XCL_BO_SYNC_BO_FROM_DEVICE, kLogitsBytes, 0); + return static_cast(logits->map()); + } +}; + +Decoder::Decoder(const Model& model, const Config& config, const std::string& kernel_dir, Transport t, int threads) + : p_(std::make_unique(model, config, kernel_dir, t, threads)) {} +Decoder::~Decoder() = default; + +const LoadStats& Decoder::load_stats() const { return p_->stats; } +const RunTimes& Decoder::last_run() const { return p_->times; } +std::string Decoder::kernels() const { return p_->ks->describe(); } + +void Decoder::feed(int token) { + if (token < 0 || token >= p_->cfg.vocab) throw std::runtime_error("token id out of range"); + const uint8_t* row = p_->embed + size_t(token) * kHidden * 2; + auto* x = static_cast(p_->xres->map()); + for (size_t i = 0; i < kHidden; ++i) { + const uint32_t u = uint32_t(row[2 * i] | (row[2 * i + 1] << 8)) << 16; + std::memcpy(x + i, &u, 4); + } + p_->write_xres(); +} + +void Decoder::set_residual(const float* v) { + std::memcpy(p_->xres->map(), v, kHidden * 4); + p_->write_xres(); +} + +void Decoder::run(size_t pos, bool head) { p_->run(pos, head); } +void Decoder::reset() { p_->reset(); } +size_t Decoder::state_bytes() const { return p_->state_bytes(); } +void Decoder::save_state(uint8_t* out) { p_->save_state(out); } +void Decoder::load_state(const uint8_t* in) { p_->load_state(in); } + +int Decoder::argmax(int n) { + const float* v = p_->read_logits(); + n = std::clamp(n, 1, int(kLogitsBytes / 4)); + int best = 0; + for (int i = 1; i < n; ++i) + if (v[i] > v[best]) best = i; + return best; +} + +void Decoder::logits(std::vector& out) { + const float* v = p_->read_logits(); + out.assign(v, v + kLogitsBytes / 4); +} + +// ---- generation --------------------------------------------------------------------------- + +int Session::pick(const GenerateOptions& opt, const std::vector& history) { + if (!(opt.repetition_penalty > 1.0f)) return dec_.argmax(vocab_); + dec_.logits(buf_); + buf_.resize(size_t(vocab_)); + const int n = int(history.size()), from = std::max(0, n - opt.penalty_window); + for (int i = from; i < n; ++i) { + const int t = history[size_t(i)]; + if (t < 0 || t >= vocab_) continue; + if (std::find(history.begin() + from, history.begin() + i, t) != history.begin() + i) continue; + float& l = buf_[size_t(t)]; + l = l > 0.0f ? l / opt.repetition_penalty : l * opt.repetition_penalty; + } + return int(std::max_element(buf_.begin(), buf_.end()) - buf_.begin()); // first maximum +} + +Session::Session(Decoder& dec, int vocab, SessionOptions o) + : dec_(dec), vocab_(vocab), turn_tokens_(std::move(o.turn_tokens)), snaps_(o.snapshots) {} + +GenerateResult Session::generate(const std::vector& prompt, const GenerateOptions& opt) { + if (prompt.empty()) throw std::runtime_error("empty prompt"); + if (prompt.size() + size_t(std::max(opt.max_tokens, 0)) > size_t(kMaxContext)) + throw std::runtime_error("prompt plus max_tokens exceeds the context (" + std::to_string(kMaxContext) + ")"); + GenerateResult res; + auto t0 = clk::now(); + // The live cache, a snapshot or nothing: whichever leaves the fewest tokens to feed + // (npu/lax_turns.h). The last prompt token runs anyway, for its logits. + // After a failed request the device may be past cache_: then only a snapshot or a reset. + const Plan plan = lax::plan(dirty_ ? std::span{} : std::span(cache_), snaps_, prompt); + res.source = plan.source; + if (plan.source == Plan::Snapshot) { + const auto t = clk::now(); + const auto& e = snaps_.at(size_t(plan.snapshot)); + dec_.load_state(e.state.data()); + cache_ = e.tokens; + snaps_.touch(size_t(plan.snapshot)); + res.restore_ms = ms_since(t); + } else if (plan.source == Plan::Scratch) { + if (dirty_ || !cache_.empty()) dec_.reset(); + cache_.clear(); + } + res.reused = cache_.size(); + dirty_ = true; // until this request completes: a throw may leave a token half run + + const auto points = snapshot_points(prompt, cache_.size(), turn_tokens_); + auto snap = points.begin(); + for (size_t i = cache_.size(); i < prompt.size(); ++i) { + if (snap != points.end() && *snap == i) { + ++snap; + if (!snaps_.contains(cache_)) { + const auto t = clk::now(); + if (auto* e = snaps_.put(cache_)) { + e->state.resize(dec_.state_bytes()); + dec_.save_state(e->state.data()); + ++res.snapshots_taken; + } + res.snapshot_ms += ms_since(t); + } + } + dec_.feed(prompt[i]); + dec_.run(i, i + 1 == prompt.size()); // the head only where a next token is wanted + cache_.push_back(prompt[i]); + } + res.prefill_ms = ms_since(t0); + + t0 = clk::now(); + std::vector history = prompt; + for (int i = 0; i < opt.max_tokens; ++i) { + const int t = pick(opt, history); + res.tokens.push_back(t); + history.push_back(t); + if (std::find(opt.stop.begin(), opt.stop.end(), t) != opt.stop.end()) { + res.stopped_at_eos = true; + break; + } + if ((opt.on_token && !opt.on_token(t)) || i + 1 == opt.max_tokens) break; + dec_.feed(t); + dec_.run(cache_.size(), true); + cache_.push_back(t); + } + dirty_ = false; + res.decode_ms = ms_since(t0); + return res; +} + +} // namespace onebit::npu::lax diff --git a/npu/lax.h b/npu/lax.h new file mode 100644 index 00000000..68d629e6 --- /dev/null +++ b/npu/lax.h @@ -0,0 +1,139 @@ +// 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 Qwen3.6-35B-A3B decode on the lax kernels (docs/npu-lax.md), in C++: every weight +// packed at load time from model.q4nx straight into the XRT buffers (npu/lax_pack.h), the +// 40 layers of a token as ONE runlist submit, then the final norm and the lm head. +// +// Per token: the residual `xres` (f32 [2048]) is the token's embedding; the full-attention +// stream is patched for the position; the runlist runs the 40 layers (each layer's +// arguments: pool xres consts kv|dkv act ptab state|dstate cfg; a linear layer gets the +// dummy kv, a full one the dummy state); `ln` normalizes xres into hn and `lm` writes +// f32 logits over the padded 248320 rows. +#pragma once + +#include "lax_kernels.h" +#include "lax_pack.h" +#include "lax_turns.h" +#include "model.h" + +#include +#include +#include +#include +#include +#include + +namespace onebit::npu::lax { + +struct LoadStats { + double buffers_ms = 0, pack_ms = 0, total_ms = 0; + size_t packed_bytes = 0; +}; + +// The last run(): host preparation before the submit (position, runlist), the 40-layer +// runlist (submit to completion), the norm and the lm head; ahead_ms is the part of the +// runlist's time the host spent preparing the next position (full ELFs). +struct RunTimes { + double prep_ms = 0, layers_ms = 0, head_ms = 0, ahead_ms = 0; +}; + +class Decoder { +public: + // kernel_dir as scripts/build-lax.sh's /kernels; t picks the kernel set + // (npu/lax_kernels.h): full ELFs, or the classic xclbin path for A/B. + Decoder(const Model& model, const Config& config, const std::string& kernel_dir, Transport t = Transport::Elf, + int threads = 0); + ~Decoder(); + Decoder(const Decoder&) = delete; + Decoder& operator=(const Decoder&) = delete; + + const LoadStats& load_stats() const; + const RunTimes& last_run() const; + std::string kernels() const; + + // The residual for the next run: a token's embedding, or given f32 values. + void feed(int token); + void set_residual(const float* xres); + // The 40 layers at cache position pos; with head, also the norm and the lm head. + void run(size_t pos, bool head); + // Start a new sequence at position 0: zero the DeltaNet state (the KV cache needs no + // clearing). Construction leaves the decoder in this state. + void reset(); + // The DeltaNet state (each linear layer's recurrent + conv state; state_bytes() in all), + // to host memory and back: what a position's cache holds besides the KV rows, which the + // positions after it overwrite before any reads them. After load_state, run from the + // position the state was saved at. + size_t state_bytes() const; + void save_state(uint8_t* out); + void load_state(const uint8_t* in); + // The last head's logits: argmax over the first n (first maximum wins), or all rows. + int argmax(int n); + void logits(std::vector& out); + +private: + struct Impl; + std::unique_ptr p_; +}; + +struct GenerateOptions { + int max_tokens = 256; + // As npu/generate.h: divides positive (multiplies negative) logits of the distinct ids + // among the last penalty_window of prompt + output; 1 is plain greedy. + float repetition_penalty = 1.0f; + int penalty_window = 64; + std::vector stop; // end-of-turn ids; the one generated is returned + std::function on_token; // each generated token; false stops +}; + +struct GenerateResult { + std::vector tokens; + double prefill_ms = 0, decode_ms = 0; + bool stopped_at_eos = false; + size_t reused = 0; // prompt tokens not fed again + Plan::Source source = Plan::Scratch; // where they came from + double restore_ms = 0; // loading the snapshot (part of prefill_ms) + double snapshot_ms = 0; // taking snapshots (part of prefill_ms) + int snapshots_taken = 0; +}; + +struct SessionOptions { + size_t snapshots = 4; // DeltaNet state snapshots kept (Decoder::state_bytes() each); 0: none + std::vector turn_tokens; // a snapshot is also taken before the last of these in a prompt +}; + +// One conversation on the decoder (npu/lax_turns.h): what the cache holds, and snapshots of +// the DeltaNet state at earlier prompts' prefixes. A prompt continues from the one of them +// that leaves the fewest tokens to feed; with none, it starts over (Decoder::reset). +// Greedy (first maximum over the tokenizer's vocab). +class Session { +public: + Session(Decoder& dec, int vocab, SessionOptions o = {}); + GenerateResult generate(const std::vector& prompt, const GenerateOptions& opt); + size_t cached() const { return cache_.size(); } + const Snapshots& snapshots() const { return snaps_; } + +private: + int pick(const GenerateOptions& opt, const std::vector& history); + Decoder& dec_; + int vocab_; + std::vector turn_tokens_; + Snapshots snaps_; + std::vector cache_; // the tokens at positions 0..n-1 + bool dirty_ = true; // the decoder may hold a sequence cache_ does not describe + std::vector buf_; +}; + +} // namespace onebit::npu::lax diff --git a/npu/lax_elf.cpp b/npu/lax_elf.cpp new file mode 100644 index 00000000..8707f2f1 --- /dev/null +++ b/npu/lax_elf.cpp @@ -0,0 +1,89 @@ +// 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 "lax_elf.h" + +#include "lax_pack.h" + +#include +#include +#include + +namespace onebit::npu::lax { + +std::string pdi_path(const std::string& kind_dir) { + for (const char* f : {"/main.pdi", "/final.prj/main.pdi"}) + if (std::filesystem::exists(kind_dir + f)) return kind_dir + f; + throw std::runtime_error(kind_dir + ": no main.pdi (scripts/build-lax.sh exports it)"); +} + +Bytes unfold(Bytes e) { + constexpr uint32_t kFold = 0x80000000u; + constexpr uint32_t kTranslatedArgs = 5; // the classic firmware translates these itself + const ElfSection ct = elf_section(e, ".ctrltext"), rela = elf_section(e, ".rela.dyn"); + auto word = [&](size_t o) { + uint32_t v; + std::memcpy(&v, e.data() + o, 4); + return v; + }; + auto clear = [&](size_t o) { + const uint32_t v = word(o) & ~kFold; + std::memcpy(e.data() + o, &v, 4); + }; + for (size_t i = 4; 4 * i + 4 <= ct.size;) { + const size_t o = ct.offset + 4 * i; + const uint32_t op = word(o); + if (op == 0x81 && 4 * (i + 11) < ct.size && (word(o + 40) & kFold)) { + if (word(o + 32) < kTranslatedArgs) + throw std::runtime_error("unfold: a DDR_PATCH on argument " + std::to_string(word(o + 32)) + + " carries bit 31, which is not the fold"); + clear(o + 40); + } + i += op_words(op); + } + for (size_t k = 0; k < rela.size / 12; ++k) + if (word(rela.offset + 12 * k + 8) & kFold) clear(rela.offset + 12 * k + 8); + refresh_uid(e); + return e; +} + +ElfInputs ElfInputs::read(const std::string& dir) { + ElfInputs in; + in.lax_l = unfold(read_file(dir + "/lax_l/insts.elf")); + in.lax_a = read_file(dir + "/lax_a/insts.elf"); // unfolded per position, after the sites are written + in.pdi = read_file(pdi_path(dir + "/lax_l")); + in.ln = unfold(read_file(dir + "/ln/insts.elf")); + in.ln_pdi = read_file(pdi_path(dir + "/ln")); + in.lm = unfold(read_file(dir + "/lm_head_q8/insts.elf")); + in.lm_pdi = read_file(pdi_path(dir + "/lm_head_q8")); + in.sites = elf_position_sites(in.lax_a, kKvRow); + return in; +} + +Bytes init_elf(const ElfInputs& in) { return assemble_full_elf(in.lax_l, in.pdi, kInitKernel, PdiMode::kInitOnly); } +Bytes linear_elf(const ElfInputs& in) { return assemble_full_elf(in.lax_l, in.pdi, kLinearKernel, PdiMode::kNone); } + +std::string full_kernel_name(size_t pos) { return "axf" + std::to_string(pos); } + +Bytes full_elf(const ElfInputs& in, size_t pos, bool with_pdi) { + if (pos >= size_t(kMaxContext)) throw std::runtime_error("position " + std::to_string(pos) + " past the context"); + Bytes e = in.lax_a; + apply_elf_position(e, in.sites, pos, kKvRow, kPtabRow); + return assemble_full_elf(unfold(std::move(e)), with_pdi ? in.pdi : Bytes{}, full_kernel_name(pos), PdiMode::kNone); +} + +Bytes norm_elf(const ElfInputs& in) { return assemble_full_elf(in.ln, in.ln_pdi, kNormKernel, PdiMode::kLoad); } +Bytes head_elf(const ElfInputs& in) { return assemble_full_elf(in.lm, in.lm_pdi, kHeadKernel, PdiMode::kLoad); } + +} // namespace onebit::npu::lax diff --git a/npu/lax_elf.h b/npu/lax_elf.h new file mode 100644 index 00000000..a26523f9 --- /dev/null +++ b/npu/lax_elf.h @@ -0,0 +1,82 @@ +// 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 lax decode's full ELFs (docs/npu-lax.md, "Full ELFs"), assembled in memory from the +// design's instruction ELFs and PDIs by npu/full_elf. Pure host code, no XRT: the kernel +// set (npu/lax_elf_kernels.cpp) opens what this builds, and the tests compare it byte for +// byte with the ELFs the reference harness ran. +// +// A kernel directory (scripts/build-lax.sh's /kernels) holds per kind +// /insts.elf and /main.pdi (or an aiecc build tree's final.prj/main.pdi): +// lax_l the merged whole-layer design's linear-attention text; its PDI configures +// the array for both layer texts +// lax_a the full-attention text (its PDI differs from lax_l's only in the build ID) +// ln the final RMSNorm, stand-alone +// lm_head_q8 the q8 lm head, stand-alone +// +// Kernels (the names XRT looks them up by): +// laxinit lax_l's PDI and only load_pdi, one unreferenced buffer argument; the layer +// context is created from it, and it heads every token's runlist +// lxf lax_l's text, no load_pdi (a config added to the layer context) +// axf lax_a's text written for position pos, no load_pdi (and no PDI) +// ln, lm load_pdi + text, each its own context +// Instruction ELFs from the default (folded) aiecc build carry +0x80000000 in the +// addends and DDR_PATCH words of arguments >= 5 (mlir-aie's kDDRAIEAddrOffset, for the +// classic firmware, which translates only the first 5 arguments); XRT's ELF flow adds +// the aperture offset to every argument itself, so the fold would land twice. unfold() +// clears it (assemble_full_elf clears bit 31 as well) and refreshes the UID note to the +// md5 of the control text XRT will run, as the reference harness does. +#pragma once + +#include "full_elf.h" +#include "lax_stream.h" + +#include +#include + +namespace onebit::npu::lax { + +struct ElfInputs { + Bytes lax_l, lax_a, pdi; // pdi: lax_l's + Bytes ln, ln_pdi, lm, lm_pdi; + std::vector sites; // lax_a's position sites + + // Reads dir//insts.elf and the PDIs; derives the position sites. + static ElfInputs read(const std::string& dir); +}; + +// The instruction ELF without the fold, its UID note refreshed. Throws on a DDR_PATCH word +// with bit 31 on an argument below 5 (not a fold). +Bytes unfold(Bytes insts_elf); + +// A kind's PDI path: /main.pdi, else /final.prj/main.pdi. +std::string pdi_path(const std::string& kind_dir); + +Bytes init_elf(const ElfInputs& in); +Bytes linear_elf(const ElfInputs& in); +// with_pdi false: an empty .pdi.1. A position config has no load_pdi, so XRT never uploads +// its PDI, but it keeps every added ELF in memory: with lax_l's 257 KiB PDI in each, the +// configs would hold about 1 GB at 4096 positions; without it they hold a few MB. +Bytes full_elf(const ElfInputs& in, size_t pos, bool with_pdi = false); +Bytes norm_elf(const ElfInputs& in); +Bytes head_elf(const ElfInputs& in); +std::string full_kernel_name(size_t pos); + +inline constexpr const char* kInitKernel = "laxinit"; +inline constexpr const char* kLinearKernel = "lxf"; +inline constexpr const char* kNormKernel = "ln"; +inline constexpr const char* kHeadKernel = "lm"; + +} // namespace onebit::npu::lax diff --git a/npu/lax_elf_kernels.cpp b/npu/lax_elf_kernels.cpp new file mode 100644 index 00000000..9ad7a21a --- /dev/null +++ b/npu/lax_elf_kernels.cpp @@ -0,0 +1,133 @@ +// 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 full-ELF kernel set of the lax decode (npu/lax_kernels.h, docs/npu-lax.md): every +// kernel is PDI + control code, assembled in memory (npu/lax_elf.h) and opened through +// XRT's ELF flow. No xclbin is read. +// +// Layer context: created from laxinit; lxf and one axf per full-attention position +// reached are configs added to it. Every token's runlist starts with the laxinit run: ln +// and lm_head run in their own stand-alone contexts, and after them the layer context's +// array configuration is gone (an init-once runlist timed out at token 1). +// +// The position lives in the relocation addends on the ELF path, so each position is its +// own config (npu/lax_stream.h, elf_position_sites), added the first time the decode +// reaches it and kept for later sequences. +#include "lax_elf.h" +#include "lax_kernels.h" +#include "lax_pack.h" + +#include +#include + +#include +#include +#include +#include + +namespace onebit::npu::lax { + +namespace { + +xrt::elf as_elf(const Bytes& b) { return xrt::elf(b.data(), b.size()); } + +class FullElf final : public KernelSet { +public: + FullElf(xrt::device& dev, const std::string& dir) : dir_(dir), in_(ElfInputs::read(dir)) { + layer_ctx_ = xrt::hw_context(dev, as_elf(init_elf(in_))); + init_ = xrt::ext::kernel(layer_ctx_, kInitKernel); + layer_ctx_.add_config(as_elf(linear_elf(in_))); + lxf_ = xrt::ext::kernel(layer_ctx_, kLinearKernel); + ln_ctx_ = xrt::hw_context(dev, as_elf(norm_elf(in_))); + ln_ = xrt::ext::kernel(ln_ctx_, kNormKernel); + lm_ctx_ = xrt::hw_context(dev, as_elf(head_elf(in_))); + lm_ = xrt::ext::kernel(lm_ctx_, kHeadKernel); + // The stand-alone kernels are built; only the full-attention text is needed again. + in_.ln.clear(), in_.ln_pdi.clear(), in_.lm.clear(), in_.lm_pdi.clear(), in_.lax_l.clear(); + + dummy_ = xrt::ext::bo(dev, 4096); + head_ = xrt::run(init_); + head_.set_arg(0, dummy_); + } + + xrt::hw_context& layer_context() override { return layer_ctx_; } + + xrt::run make_run(Stage s) override { + switch (s) { + case Stage::LayerLinear: return xrt::run(lxf_); + case Stage::LayerFull: return xrt::run(axf_kernel(pos_)); + case Stage::Norm: return xrt::run(ln_); + case Stage::Head: return xrt::run(lm_); + } + throw std::runtime_error("unknown stage"); + } + + int first_buffer_arg() const override { return 0; } + + void set_position(size_t pos) override { + axf_kernel(pos); + pos_ = pos; + } + + bool position_in_run() const override { return true; } + + std::optional head_run() override { return head_; } + + std::string describe() const override { + char b[160]; + std::snprintf(b, sizeof b, "; %zu position configs, %.2f ms each to add", axf_.size(), + axf_.empty() ? 0.0 : config_ms_ / double(axf_.size())); + return "full ELFs from " + dir_ + b; + } + +private: + // The full-attention config for position pos, added to the layer context on first use. + xrt::kernel& axf_kernel(size_t pos) { + if (auto it = axf_.find(pos); it != axf_.end()) return it->second; + const auto t0 = std::chrono::steady_clock::now(); + layer_ctx_.add_config(as_elf(full_elf(in_, pos))); + auto& k = axf_.emplace(pos, xrt::ext::kernel(layer_ctx_, full_kernel_name(pos))).first->second; + config_ms_ += std::chrono::duration(std::chrono::steady_clock::now() - t0).count(); + return k; + } + + std::string dir_; + ElfInputs in_; + xrt::hw_context layer_ctx_, ln_ctx_, lm_ctx_; + xrt::kernel init_, lxf_, ln_, lm_; + std::map axf_; + xrt::bo dummy_; + xrt::run head_; + size_t pos_ = 0; + double config_ms_ = 0; +}; + +} // namespace + +std::unique_ptr elf_kernels(xrt::device& dev, const std::string& dir) { + return std::make_unique(dev, dir); +} + +Transport parse_transport(const std::string& name) { + if (name == "elf") return Transport::Elf; + if (name == "classic") return Transport::Classic; + throw std::runtime_error("kernel transport '" + name + "': elf or classic"); +} + +std::unique_ptr make_kernels(xrt::device& dev, const std::string& dir, Transport t) { + return t == Transport::Elf ? elf_kernels(dev, dir) : classic_kernels(dev, dir); +} + +} // namespace onebit::npu::lax diff --git a/npu/lax_kernels.h b/npu/lax_kernels.h new file mode 100644 index 00000000..7a333473 --- /dev/null +++ b/npu/lax_kernels.h @@ -0,0 +1,82 @@ +// 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. + +// How the lax decode (npu/lax.h) gets its four kernels onto the device. The decoder only +// asks for runs with their non-buffer arguments set, the context the layer runlist is +// bound to, and a way to set the full-attention position; everything about xclbins, +// instruction buffers or ELFs stays behind this interface. +// +// Two sets: +// elf_kernels full ELFs (docs/npu-lax.md, "Full ELFs"): PDI + control code, +// assembled in memory by npu/full_elf; no xclbin. The default. +// classic_kernels xclbin + insts.bin, the path the design was brought up on; kept +// selectable for A/B. +#pragma once + +#include +#include +#include +#include + +#include +#include +#include + +namespace onebit::npu::lax { + +enum class Stage { LayerLinear, LayerFull, Norm, Head }; + +class KernelSet { +public: + virtual ~KernelSet() = default; + // The context both layer kernels live in: one runlist carries all 40 layers. + virtual xrt::hw_context& layer_context() = 0; + // A run of that stage with every argument before the buffers set. A LayerFull run is + // for the position set last. + virtual xrt::run make_run(Stage s) = 0; + // The argument index of the stage's first buffer. + virtual int first_buffer_arg() const = 0; + // Point the full-attention layers at cache position pos (their KV window, new row and + // RoPE record); called before every token's runlist. + virtual void set_position(size_t pos) = 0; + // True: the position is baked into each LayerFull run (a full ELF config per + // position); set_position only affects runs made after it, so the decoder makes new + // ones, and may do so while an earlier runlist executes. False: set_position rewrites + // what the existing runs execute (classic: the shared instruction buffer). + virtual bool position_in_run() const = 0; + // A run that must head every token's runlist, or none. Full ELFs: the init config's + // load_pdi, because the norm and the lm head run in their own contexts, and a full-ELF + // context does not get its array configuration back after another context ran. + virtual std::optional head_run() = 0; + virtual std::string describe() const = 0; +}; + +// The classic path: dir/{lax_l,lax_a,ln,lm_head_q8}/{final.xclbin,insts.bin}, as +// scripts/build-lax.sh builds them. lax_l and lax_a share one xclbin (one context); the +// instruction words go in through an instruction buffer at argument 1. +std::unique_ptr classic_kernels(xrt::device& dev, const std::string& dir); + +// The full-ELF path: dir/{lax_l,lax_a,ln,lm_head_q8}/{insts.elf,main.pdi}, as +// scripts/build-lax.sh exports them (docs/npu-lax.md, "Kernel directory"). The layer +// context is created from an init config (lax_l's PDI, load_pdi only); the linear layer +// and each full-attention position reached are configs added to it; the norm and the lm +// head are stand-alone contexts. Buffers from argument 0. +std::unique_ptr elf_kernels(xrt::device& dev, const std::string& dir); + +enum class Transport { Elf, Classic }; +Transport parse_transport(const std::string& name); // "elf" or "classic" +std::unique_ptr make_kernels(xrt::device& dev, const std::string& dir, Transport t); + +} // namespace onebit::npu::lax diff --git a/npu/lax_pack.cpp b/npu/lax_pack.cpp new file mode 100644 index 00000000..f486d5df --- /dev/null +++ b/npu/lax_pack.cpp @@ -0,0 +1,443 @@ +// 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. + +// Built without floating-point contraction (CMakeLists.txt): requant_q8 must round every +// float32 step on its own, as the reference packer's NumPy does. +#include "lax_pack.h" + +#include "q4nx.h" + +#include + +#include +#include +#include +#include +#include + +namespace onebit::npu::lax { + +namespace { + +float f32_of_bf16(uint16_t v) { + const uint32_t u = uint32_t(v) << 16; + float f; + std::memcpy(&f, &u, 4); + return f; +} + +uint32_t bits(float f) { + uint32_t u; + std::memcpy(&u, &f, 4); + return u; +} + +// f32 -> bf16 toward -inf, toward +inf, and to nearest even. +uint16_t bf16_floor(float x) { + const uint32_t u = bits(x); + return uint16_t((u >> 31) ? (u + 0xFFFFu) >> 16 : u >> 16); +} +uint16_t bf16_ceil(float x) { + const uint32_t u = bits(x); + return uint16_t((u >> 31) ? u >> 16 : (u + 0xFFFFu) >> 16); +} +uint16_t bf16_rne(float x) { + const uint32_t u = bits(x); + return uint16_t((u + 0x7FFFu + ((u >> 16) & 1u)) >> 16); +} + +uint16_t rd16(const uint8_t* p) { return uint16_t(p[0] | (p[1] << 8)); } +void wr16(uint8_t* p, uint16_t v) { + p[0] = uint8_t(v); + p[1] = uint8_t(v >> 8); +} + +// Nibble / int8 code index of (row r, 32-column group g, lane i) inside a chunk. +constexpr size_t code_index(int r, int g, int i) { return size_t((r / 16) * 4096 + g * 512 + i * 16 + r % 16); } + +std::string lname(int layer, const char* suffix) { return "model.layer." + std::to_string(layer) + "." + suffix; } + +size_t chunk_bytes(const Tensor& t) { + if (t.dtype != "I8" || t.shape.empty()) throw std::runtime_error("not a quantized tensor"); + return size_t(t.shape.back()); +} + +// Chunks [c0, c0 + n) of a quantized tensor as q4_1: native q4_1 in place, anything else +// converted into `store`. Mirrors the reference packer's source handling: q8 is +// re-quantized, Q4_K transcoded, and a 5120-byte tensor whose first selected chunk has +// every min zero is the signed-nibble form. +class Q41 { +public: + Q41(const Tensor& t, const std::string& name, size_t c0, size_t n) { + const size_t cb = chunk_bytes(t); + if (cb != kQ41 && cb != kQ8 && cb != kQ4K) + throw std::runtime_error(name + ": " + std::to_string(cb) + "-byte chunks; the packer reads 5120, 8704 and 4736"); + if ((c0 + n) * cb > t.bytes) throw std::runtime_error(name + ": too few chunks"); + const uint8_t* src = t.data + c0 * cb; + if (cb == kQ41 && !is_signed_q4(src)) { + base_ = src; + return; + } + store_.resize(n * kQ41); + for (size_t i = 0; i < n; ++i) { + uint8_t* d = store_.data() + i * kQ41; + if (cb == kQ8) requant_q8(src + i * cb, d); + else if (cb == kQ4K) q4nx::q4k_to_q4_1(src + i * cb, d); + else q4_0_to_q4_1(src + i * cb, d); + } + base_ = store_.data(); + } + const uint8_t* operator[](size_t i) const { return base_ + i * kQ41; } + +private: + const uint8_t* base_ = nullptr; + std::vector store_; +}; + +// std band law: nch pool chunks of an [out, in_dim] tensor starting at source chunk c0. +void std_band(const Model& m, const std::string& name, size_t c0, size_t nch, size_t in_dim, uint8_t* dst) { + if (in_dim % 256) throw std::runtime_error(name + ": in_dim is not whole 256-column tiles"); + const Tensor& t = m.tensor(name); + const Q41 src(t, name, c0, nch); + for (size_t c = 0; c < nch; ++c) std::memcpy(dst + c * kQ41, src[std_src(c, in_dim)], kQ41); +} + +void put(const Model& m, const std::string& name, size_t cap, uint8_t* dst) { + const Tensor& t = m.tensor(name); + if (t.bytes > cap) throw std::runtime_error(name + ": does not fit its " + std::to_string(cap) + " B slot"); + std::memcpy(dst, t.data, t.bytes); +} + +int need(const nlohmann::json& j, const char* key) { + if (!j.contains(key)) throw std::runtime_error(std::string("config.json has no ") + key); + return j.at(key).get(); +} + +} // namespace + +// ---- the chunk transforms -------------------------------------------------------------- + +void requant_q8(const uint8_t* q8, uint8_t* q41) { + std::memset(q41, 0, kQ41); + for (int r = 0; r < 32; ++r) { + for (int g = 0; g < 8; ++g) { + const int meta = g * 32 + r; + const float sc = f32_of_bf16(rd16(q8 + 2 * meta)); + float v[32]; + float mn = 0, mx = 0; + for (int i = 0; i < 32; ++i) { + v[i] = float(int8_t(q8[512 + code_index(r, g, i)])) * sc; // exact: 8 x 8 significant bits + // Ties go to the later element (x86 minps / maxps, as NumPy reduces): in a + // block whose scale is zero every value is +0 or -0, and the min's sign is + // stored in m. + mn = i && mn < v[i] ? mn : v[i]; + mx = i && mx > v[i] ? mx : v[i]; + } + const uint16_t mu = bf16_floor(mn); + const float mf = f32_of_bf16(mu); + const float range = mx - mf; + const uint16_t du = bf16_ceil(range / 15.0f); + const float d = f32_of_bf16(du); + const float inv = d > 0.0f ? 1.0f / d : 0.0f; + wr16(q41 + 2 * meta, du); + wr16(q41 + 512 + 2 * meta, mu); + for (int i = 0; i < 32; ++i) { + const float diff = v[i] - mf; + const float scaled = diff * inv; + const float t = scaled + 0.5f; + const int q = std::clamp(int(t), 0, 15); // truncation toward zero, then clip + const size_t n = code_index(r, g, i); + q41[1024 + n / 2] |= uint8_t((n & 1) ? q << 4 : q); + } + } + } +} + +bool is_signed_q4(const uint8_t* c) { + for (size_t i = 512; i < 1024; ++i) + if (c[i]) return false; + return true; +} + +void q4_0_to_q4_1(const uint8_t* s, uint8_t* d) { + std::memcpy(d, s, 512); + for (int k = 0; k < 256; ++k) wr16(d + 512 + 2 * k, bf16_rne(-8.0f * f32_of_bf16(rd16(s + 2 * k)))); + for (size_t i = 1024; i < kQ41; ++i) d[i] = uint8_t(s[i] ^ 0x88); +} + +size_t std_src(size_t c, size_t in_dim) { + // A band is 64 rows x in_dim: in_dim/128 chunks, chunk i of a band covers row half i%2 + // and k-tile i/2. The source is a plain raster of 32-row blocks x 256-column tiles. + const size_t per_band = in_dim / 128; + const size_t rb = 2 * (c / per_band) + c % 2, kt = (c % per_band) / 2; + return rb * (in_dim / 256) + kt; +} + +size_t down_src(size_t c) { return 2 * (4 * (c / 8) + c % 4) + (c / 4) % 2; } + +size_t stripe_src(size_t c) { return 8 * (c % 4) + c / 4; } // in_dim 2048: 8 k-tiles per row block + +size_t lmhead_src(size_t k, size_t hidden) { + // 128-row supertiles: pool chunk c of a band is (row quarter c%4, k-tile c/4). + const size_t nk = hidden / 256, per_band = 4 * nk; + const size_t s = k / per_band, r = k % per_band; + return (4 * s + r % 4) * nk + r / 4; +} + +// ---- configuration --------------------------------------------------------------------- + +Config Config::from_model(const Model& m, const std::string& dir) { + std::ifstream f(dir + "/config.json"); + if (!f) throw std::runtime_error("cannot read " + dir + "/config.json"); + const auto j = nlohmann::json::parse(f); + const auto& t = j.contains("text_config") ? j.at("text_config") : j; + Config c; + c.hidden = need(t, "hidden_size"); + c.layers = need(t, "num_hidden_layers"); + c.heads = need(t, "num_attention_heads"); + c.kv_heads = need(t, "num_key_value_heads"); + c.head_dim = need(t, "head_dim"); + c.experts = need(t, "num_experts"); + c.topk = need(t, "num_experts_per_tok"); + c.moe_inter = need(t, "moe_intermediate_size"); + c.shared_inter = need(t, "shared_expert_intermediate_size"); + c.lin_key_heads = need(t, "linear_num_key_heads"); + c.lin_value_heads = need(t, "linear_num_value_heads"); + c.lin_key_dim = need(t, "linear_key_head_dim"); + c.lin_value_dim = need(t, "linear_value_head_dim"); + c.conv_kernel = need(t, "linear_conv_kernel_dim"); + const auto& rp = t.contains("rope_parameters") ? t.at("rope_parameters") : t; + c.rope_theta = rp.value("rope_theta", t.value("rope_theta", 0.0)); + const double prf = rp.value("partial_rotary_factor", t.value("partial_rotary_factor", 1.0)); + c.rotary_dim = int(c.head_dim * prf); + for (const auto& lt : t.at("layer_types")) { + const auto s = lt.get(); + if (s == "linear_attention") c.kinds.push_back(Kind::Linear); + else if (s == "full_attention") c.kinds.push_back(Kind::Full); + else throw std::runtime_error("layer type " + s + " has no lax kernel"); + } + c.vocab = int(m.tensor("model.embed_tokens.weight").shape.at(0)); + + // The kernels are compiled for exactly this geometry (recipes/specs/qwen36-35b-a3b.json). + const bool ok = c.hidden == 2048 && c.layers == 40 && int(c.kinds.size()) == 40 && c.heads == 16 && + c.kv_heads == 2 && c.head_dim == 256 && c.rotary_dim == 64 && c.experts == 256 && c.topk == 8 && + c.moe_inter == 512 && c.shared_inter == 512 && c.lin_key_heads == 16 && c.lin_value_heads == 32 && + c.lin_key_dim == 128 && c.lin_value_dim == 128 && c.conv_kernel == 4 && c.vocab == 248320 && + c.rope_theta > 0; + if (!ok) throw std::runtime_error(dir + ": not the Qwen3.6-35B-A3B geometry the lax kernels are built for"); + return c; +} + +size_t consts_bytes(Kind k) { return k == Kind::Linear ? kConstsBytes : kConstsFull; } + +// ---- the buffers ---------------------------------------------------------------------- + +void pack_pool(const Model& m, const Config& c, int L, uint8_t* dst) { + std::memset(dst, 0, kPoolBytes); + const MoeLayout g; + const size_t E = size_t(c.experts), stripes = size_t(c.moe_inter) / 128, per_stripe = g.stripe / kQ41; + + // Routed experts' up and gate: per expert, `stripes` pairs [up stripe k | gate stripe k], + // each stripe's 4 row blocks x 8 k-tiles transposed to k-tile-major. + { + const std::string un = lname(L, "mlp.up_exps_proj.weight"), gn = lname(L, "mlp.gate_exps_proj.weight"); + const Tensor &ut = m.tensor(un), > = m.tensor(gn); + const size_t n = E * stripes * per_stripe; + const Q41 up(ut, un, 0, n), gate(gt, gn, 0, n); + for (size_t e = 0; e < E; ++e) + for (size_t k = 0; k < stripes; ++k) { + const size_t src = (stripes * e + k) * per_stripe; + uint8_t* d = dst + (2 * stripes * e + 2 * k) * g.stripe; + for (size_t j = 0; j < per_stripe; ++j) { + std::memcpy(d + j * kQ41, up[src + stripe_src(j)], kQ41); + std::memcpy(d + g.stripe + j * kQ41, gate[src + stripe_src(j)], kQ41); + } + } + } + // Routed experts' down slices. + { + const std::string dn = lname(L, "mlp.down_exps_proj.weight"); + const size_t per = g.up_bytes / kQ41; + const Q41 down(m.tensor(dn), dn, 0, E * per); + for (size_t e = 0; e < E; ++e) + for (size_t j = 0; j < per; ++j) + std::memcpy(dst + g.pool_down + e * g.up_bytes + j * kQ41, down[e * per + down_src(j)], kQ41); + } + // The shared expert. + const size_t H = size_t(c.hidden), FF = size_t(c.shared_inter), band = FF * H / 8192; + std_band(m, lname(L, "mlp.share_up_exps_proj.weight"), 0, band, H, dst + g.share_up); + std_band(m, lname(L, "mlp.share_gate_exps_proj.weight"), 0, band, H, dst + g.share_gate); + std_band(m, lname(L, "mlp.share_down_exps_proj.weight"), 0, band, FF, dst + g.share_down); + + uint8_t* a = dst + g.attn; + if (c.kinds[size_t(L)] == Kind::Linear) { + // DeltaNet: qkv (2*16*128 + 32*128 = 8192 rows), then the z gate (4096 rows). + const size_t qkv = size_t(2 * c.lin_key_heads * c.lin_key_dim + c.lin_value_heads * c.lin_value_dim) * H / 8192; + const size_t z = size_t(c.lin_value_heads * c.lin_value_dim) * H / 8192; + std_band(m, lname(L, "linear_attn.qkv_proj.weight"), 0, qkv, H, a); + std_band(m, lname(L, "self_attn.gate_proj.weight"), 0, z, H, a + qkv * kQ41); + } else { + // Full attention: q_proj's first half, k, v, q_proj's second half, o. + const size_t qh = size_t(c.heads * c.head_dim) * H / 8192, kv = size_t(c.kv_heads * c.head_dim) * H / 8192; + const size_t o = H * size_t(c.heads * c.head_dim) / 8192; + const std::string q = lname(L, "self_attn.q_proj.weight"); + std_band(m, q, 0, qh, H, a); + a += qh * kQ41; + std_band(m, lname(L, "self_attn.k_proj.weight"), 0, kv, H, a); + a += kv * kQ41; + std_band(m, lname(L, "self_attn.v_proj.weight"), 0, kv, H, a); + a += kv * kQ41; + std_band(m, q, qh, qh, H, a); + a += qh * kQ41; + std_band(m, lname(L, "self_attn.o_proj.weight"), 0, o, size_t(c.heads * c.head_dim), a); + } +} + +void pack_consts(const Model& m, const Config& c, int L, uint8_t* dst) { + const Kind k = c.kinds[size_t(L)]; + std::memset(dst, 0, consts_bytes(k)); + const size_t H = size_t(c.hidden); + if (k == Kind::Full) { + put(m, lname(L, "input_layernorm.weight"), 4096, dst + 0); + put(m, lname(L, "post_attention_layernorm.weight"), 4096, dst + 4096); + put(m, lname(L, "self_attn.q_norm.weight"), 512, dst + 8192); + put(m, lname(L, "self_attn.k_norm.weight"), 512, dst + 8704); + put(m, lname(L, "moe_router.weight"), 1u << 20, dst + 10240); + put(m, lname(L, "shared_expert_gate.weight"), 4096, dst + 1058816); + return; + } + put(m, lname(L, "input_layernorm.weight"), 4096, dst + 0); + put(m, lname(L, "linear_attn.ssm_alpha_proj.weight"), 131072, dst + 4096); + put(m, lname(L, "linear_attn.ssm_beta_proj.weight"), 131072, dst + 135168); + put(m, lname(L, "linear_attn.ssm_a"), 128, dst + 266240); + put(m, lname(L, "linear_attn.ssm_dt.bias"), 128, dst + 266368); + { + // conv1d bf16 [taps, 8 x 1024] -> [8 groups][taps][1024] + const std::string n = lname(L, "linear_attn.ssm_conv1d.weight"); + const Tensor& t = m.tensor(n); + const size_t taps = size_t(c.conv_kernel), groups = 8, width = 1024; + if (t.bytes != taps * groups * width * 2) throw std::runtime_error(n + ": unexpected size"); + uint8_t* d = dst + 270336; + for (size_t gi = 0; gi < groups; ++gi) + for (size_t ti = 0; ti < taps; ++ti) + std::memcpy(d + (gi * taps + ti) * width * 2, t.data + (ti * groups + gi) * width * 2, width * 2); + } + put(m, lname(L, "linear_attn.ssm_norm.weight"), 256, dst + 335872); + put(m, lname(L, "post_attention_layernorm.weight"), 4096, dst + 339968); + put(m, lname(L, "moe_router.weight"), 1u << 20, dst + 344064); + put(m, lname(L, "shared_expert_gate.weight"), 4096, dst + 1392640); + const size_t vd = size_t(c.lin_value_heads * c.lin_value_dim); + std_band(m, lname(L, "linear_attn.ssm_out_proj.weight"), 0, H * vd / 8192, vd, dst + 1396736); +} + +void pack_lmhead(const Model& m, uint8_t* dst) { + const Tensor& t = m.tensor("lm_head.weight"); + if (chunk_bytes(t) != kQ8) throw std::runtime_error("lm_head.weight: the lax lm head is q8"); + const size_t n = t.bytes / kQ8; + if (n * kQ8 > kLmPoolBytes) throw std::runtime_error("lm_head larger than its pool"); + for (size_t k = 0; k < n; ++k) std::memcpy(dst + k * kQ8, t.data + lmhead_src(k, kHidden) * kQ8, kQ8); + std::memset(dst + n * kQ8, 0, kLmPoolBytes - n * kQ8); +} + +void pack_norm(const Model& m, uint8_t* dst) { put(m, "model.norm.weight", kHidden * 2, dst); } + +void pack_ptab(const Config& c, int rows, uint8_t* dst) { + // Row p: i32 valid = p (the cached rows before it), i32 nf = max(p, 1) (rows streamed: + // position 0 streams one masked dummy), then the RoPE pairs (i, i + rot/2) as f32 cos at + // 512 and f32 sin right after. inv_freq = theta^(-i/half) and the angle are double, the + // stored values float32, as the reference computes them. + const int half = c.rotary_dim / 2; + std::vector inv(static_cast(half)); + for (int i = 0; i < half; ++i) inv[size_t(i)] = std::pow(c.rope_theta, -double(i) / double(half)); + std::memset(dst, 0, size_t(rows) * kPtabRow); + for (int p = 0; p < rows; ++p) { + uint8_t* row = dst + size_t(p) * kPtabRow; + const int32_t valid = p, nf = std::max(p, 1); + std::memcpy(row, &valid, 4); + std::memcpy(row + 4, &nf, 4); + for (int i = 0; i < half; ++i) { + const double ang = double(p) * inv[size_t(i)]; + const float cs = float(std::cos(ang)), sn = float(std::sin(ang)); + std::memcpy(row + 512 + 4 * i, &cs, 4); + std::memcpy(row + 512 + 4 * half + 4 * i, &sn, 4); + } + } +} + +// ---- SHA-256 ------------------------------------------------------------------------- + +std::array sha256(std::span data) { + static constexpr uint32_t K[64] = { + 0x428a2f98, 0x71374491, 0xb5c0fbcf, 0xe9b5dba5, 0x3956c25b, 0x59f111f1, 0x923f82a4, 0xab1c5ed5, + 0xd807aa98, 0x12835b01, 0x243185be, 0x550c7dc3, 0x72be5d74, 0x80deb1fe, 0x9bdc06a7, 0xc19bf174, + 0xe49b69c1, 0xefbe4786, 0x0fc19dc6, 0x240ca1cc, 0x2de92c6f, 0x4a7484aa, 0x5cb0a9dc, 0x76f988da, + 0x983e5152, 0xa831c66d, 0xb00327c8, 0xbf597fc7, 0xc6e00bf3, 0xd5a79147, 0x06ca6351, 0x14292967, + 0x27b70a85, 0x2e1b2138, 0x4d2c6dfc, 0x53380d13, 0x650a7354, 0x766a0abb, 0x81c2c92e, 0x92722c85, + 0xa2bfe8a1, 0xa81a664b, 0xc24b8b70, 0xc76c51a3, 0xd192e819, 0xd6990624, 0xf40e3585, 0x106aa070, + 0x19a4c116, 0x1e376c08, 0x2748774c, 0x34b0bcb5, 0x391c0cb3, 0x4ed8aa4a, 0x5b9cca4f, 0x682e6ff3, + 0x748f82ee, 0x78a5636f, 0x84c87814, 0x8cc70208, 0x90befffa, 0xa4506ceb, 0xbef9a3f7, 0xc67178f2}; + uint32_t h[8] = {0x6a09e667, 0xbb67ae85, 0x3c6ef372, 0xa54ff53a, 0x510e527f, 0x9b05688c, 0x1f83d9ab, 0x5be0cd19}; + auto rotr = [](uint32_t x, int n) { return (x >> n) | (x << (32 - n)); }; + auto block = [&](const uint8_t* p) { + uint32_t w[64]; + for (int i = 0; i < 16; ++i) + w[i] = uint32_t(p[4 * i]) << 24 | uint32_t(p[4 * i + 1]) << 16 | uint32_t(p[4 * i + 2]) << 8 | p[4 * i + 3]; + for (int i = 16; i < 64; ++i) { + const uint32_t s0 = rotr(w[i - 15], 7) ^ rotr(w[i - 15], 18) ^ (w[i - 15] >> 3); + const uint32_t s1 = rotr(w[i - 2], 17) ^ rotr(w[i - 2], 19) ^ (w[i - 2] >> 10); + w[i] = w[i - 16] + s0 + w[i - 7] + s1; + } + uint32_t a = h[0], b = h[1], c = h[2], d = h[3], e = h[4], f = h[5], g = h[6], hh = h[7]; + for (int i = 0; i < 64; ++i) { + const uint32_t t1 = hh + (rotr(e, 6) ^ rotr(e, 11) ^ rotr(e, 25)) + ((e & f) ^ (~e & g)) + K[i] + w[i]; + const uint32_t t2 = (rotr(a, 2) ^ rotr(a, 13) ^ rotr(a, 22)) + ((a & b) ^ (a & c) ^ (b & c)); + hh = g; + g = f; + f = e; + e = d + t1; + d = c; + c = b; + b = a; + a = t1 + t2; + } + h[0] += a, h[1] += b, h[2] += c, h[3] += d, h[4] += e, h[5] += f, h[6] += g, h[7] += hh; + }; + const size_t n = data.size(), full = n / 64; + for (size_t i = 0; i < full; ++i) block(data.data() + 64 * i); + uint8_t tail[128] = {}; + const size_t rem = n - full * 64; + std::memcpy(tail, data.data() + full * 64, rem); + tail[rem] = 0x80; + const size_t tl = rem < 56 ? 64 : 128; + const uint64_t bitlen = uint64_t(n) * 8; + for (int i = 0; i < 8; ++i) tail[tl - 1 - i] = uint8_t(bitlen >> (8 * i)); + for (size_t o = 0; o < tl; o += 64) block(tail + o); + std::array out{}; + for (int i = 0; i < 8; ++i) + for (int j = 0; j < 4; ++j) out[size_t(4 * i + j)] = uint8_t(h[i] >> (24 - 8 * j)); + return out; +} + +std::string hex(std::span d) { + static const char* x = "0123456789abcdef"; + std::string s; + for (uint8_t b : d) { + s += x[b >> 4]; + s += x[b & 15]; + } + return s; +} + +} // namespace onebit::npu::lax diff --git a/npu/lax_pack.h b/npu/lax_pack.h new file mode 100644 index 00000000..24d716d9 --- /dev/null +++ b/npu/lax_pack.h @@ -0,0 +1,121 @@ +// 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 buffers the Qwen3.6-35B-A3B `lax` kernels read (docs/npu-lax.md), packed from a +// Q4NX container. Pure host code. +// +// The layout is the one the open kernels' recipe (third_party/OpenFlowLM-Next, +// open_kernels/recipes/qwen36moe.py pack_plan and recipes/pack.py, MIT) produces for this +// model with every projection at q4_1 (OPEN_KERNELS_FORCE_Q4_1, make_decode.py --requant), +// which is how the lax kernels are built. It is re-implemented here from that +// specification, not copied, and checked byte for byte against the Python packer's +// output (tests/npu_lax_test.cpp, docs/npu-lax.md "The C++ driver"). +// +// Per layer the kernel takes +// pool 512 MiB: the routed experts' up/gate stripes, their down slices, the shared +// expert, then the attention projections (DeltaNet qkv + z gate, or q/k/v/o) +// consts norms, the router, the DeltaNet constants and (linear layers) ssm_out +// and the model-wide buffers are the q8 lm head pool, the final norm and the position +// table (ptab: per position [i32 valid | i32 nf | f32 cos @512 | f32 sin]). +// +// Every quantized chunk is a 32-row x 256-column tile. A q4_1 chunk (5120 B) is bf16 +// d[g*32+r] @0, bf16 m[g*32+r] @512, then nibbles: nibble (r/16)*4096 + g*512 + i*16 + r%16 +// holds column g*32+i of row r (even nibble in the low half of its byte). A q8 chunk +// (8704 B) is bf16 d[g*32+r] @0 and int8 codes @512 at the same index. +#pragma once + +#include "model.h" + +#include +#include +#include +#include +#include + +namespace onebit::npu::lax { + +constexpr size_t kQ41 = 5120, kQ8 = 8704, kQ4K = 4736; + +enum class Kind : uint8_t { Linear, Full }; // DeltaNet (lax_l) or full attention (lax_a) + +// The model dimensions the kernels were built for, read from config.json and the +// container. Anything else is refused: the kernels are compiled for these. +struct Config { + int hidden = 0, layers = 0, heads = 0, kv_heads = 0, head_dim = 0, rotary_dim = 0; + int experts = 0, topk = 0, moe_inter = 0, shared_inter = 0; + int lin_key_heads = 0, lin_value_heads = 0, lin_key_dim = 0, lin_value_dim = 0, conv_kernel = 0; + int vocab = 0; // embedding / lm head rows (248320; the tokenizer's ids stop at 248070) + double rope_theta = 0; + std::vector kinds; + + static Config from_model(const Model& m, const std::string& dir); +}; + +// Buffer sizes of the lax program (open_kernels model/lax_decode_cfg.py). +constexpr size_t kPoolBytes = 512u << 20; +constexpr size_t kConstsBytes = 11882496; // the buffer; a full-attention layer fills less +constexpr size_t kConstsFull = 1062912; +constexpr size_t kActBytes = 190464; +constexpr size_t kCfgBytes = 4096; +constexpr size_t kStateBytes = 2342912; // DeltaNet recurrent + conv state +constexpr int kMaxContext = 4096; // KV and ptab rows +constexpr size_t kKvRow = 2048; // one position's [K | V]: 2 kv heads x 256 bf16 each +constexpr size_t kKvBytes = kKvRow * kMaxContext; +constexpr size_t kPtabRow = 1024; +constexpr size_t kPtabBytes = kPtabRow * kMaxContext; +constexpr size_t kLmPoolBytes = 542113792; +constexpr size_t kHidden = 2048; +constexpr size_t kLogitsBytes = 993280; // f32 logits over the padded 248320 rows + +// Where the MoE block lies in the pool (also what the on-device router retargets). +struct MoeLayout { + size_t stripe = 163840; // 128 rows of up (or gate) x 2048 + size_t up_bytes = 655360; // one expert's up (= gate = down) + size_t pool_down = 335544320; + size_t share_up = 503316480, share_gate = 503971840, share_down = 504627200; + size_t attn = 505282560; // first byte after the shared expert +}; + +size_t consts_bytes(Kind k); + +// One layer's pool; dst holds kPoolBytes and is fully written (zeros included). +void pack_pool(const Model& m, const Config& c, int layer, uint8_t* dst); +// One layer's consts; dst holds consts_bytes(kind) and is fully written. +void pack_consts(const Model& m, const Config& c, int layer, uint8_t* dst); +// The q8 lm head pool; dst holds kLmPoolBytes. +void pack_lmhead(const Model& m, uint8_t* dst); +// The final norm, bf16 [hidden]. +void pack_norm(const Model& m, uint8_t* dst); +// rows position records of kPtabRow bytes. +void pack_ptab(const Config& c, int rows, uint8_t* dst); + +// The chunk transforms, exposed for their unit tests. +// q8 -> q4_1 per 32-value block: m = min rounded down to bf16, d = (max - m) / 15 rounded +// up, nibble = trunc((v - m) / d + 0.5) clipped to 0..15, all in float32. +void requant_q8(const uint8_t* q8, uint8_t* q41); +// The signed-nibble form of a 5120-byte chunk (every min zero) as q4_1: m = -8d, nibble ^ 8. +void q4_0_to_q4_1(const uint8_t* src, uint8_t* dst); +bool is_signed_q4(const uint8_t* first_chunk); +// Pool chunk c of a std band [out, in_dim] reads source chunk std_src(c, in_dim). +size_t std_src(size_t c, size_t in_dim); +size_t down_src(size_t c); // one expert's down slice (128 chunks) +size_t stripe_src(size_t c); // inside one 128-row up/gate stripe (32 chunks) +size_t lmhead_src(size_t k, size_t hidden); + +// SHA-256 (FIPS 180-4), for the byte-identity tables. +std::array sha256(std::span data); +std::string hex(std::span digest); + +} // namespace onebit::npu::lax diff --git a/npu/lax_stream.cpp b/npu/lax_stream.cpp new file mode 100644 index 00000000..562ac5df --- /dev/null +++ b/npu/lax_stream.cpp @@ -0,0 +1,168 @@ +// 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 "lax_stream.h" + +#include +#include +#include +#include + +namespace onebit::npu::lax { + +size_t op_words(uint32_t w) { + switch (w) { + case 0x00: return 6; + case 0x01: return 12; + case 0x03: return 7; + case 0x80: return 4; + case 0x81: return 12; + default: return 1; + } +} + +std::vector position_patches(std::span w, size_t kv_row) { + std::vector out; + size_t last_bd = 0; + bool have_bd = false; + for (size_t i = 4; i < w.size(); i += op_words(w[i])) { + if (w[i] == 0x01) { + last_bd = i; + have_bd = true; + } + if (w[i] != 0x81 || i + 11 >= w.size()) continue; + const uint32_t reg = w[i + 6], arg = w[i + 8]; + const uint32_t flags = w[i + 10] & 0x80000000u, off = w[i + 10] & 0x7FFFFFFFu; + if (arg == 3 && off == 0) { + // The fill's length is in the BD written just before: its register + 4 is the + // register this patch rewrites (the BD's address word). + if (!have_bd || last_bd + 4 >= w.size() || w[last_bd + 2] + 4 != reg) + throw std::runtime_error("lax_a stream: no BD write before the KV window fill"); + out.push_back({last_bd + 4, PosPatch::WindowLength, 0}); + out.push_back({i + 10, PosPatch::WindowOffset, flags}); + } else if (arg == 3 && off == kv_row) { + out.push_back({i + 10, PosPatch::RowDrain, flags}); + } else if (arg == 3) { + throw std::runtime_error("lax_a stream: unexpected KV transfer at offset " + std::to_string(off)); + } else if (arg == 5) { + out.push_back({i + 10, PosPatch::Record, flags}); + } + } + for (int k = 0; k < 4; ++k) { + int n = 0; + for (const auto& p : out) n += p.kind == k; + if (n != 1) + throw std::runtime_error("lax_a stream: " + std::to_string(n) + " position patches of kind " + std::to_string(k) + + ", expected 1"); + } + return out; +} + +uint32_t position_value(PosPatch::Kind kind, size_t pos, size_t kv_row, size_t ptab_row) { + const size_t nf = pos ? pos : 1; // position 0 streams one dummy row the core masks + switch (kind) { + case PosPatch::WindowLength: return uint32_t(nf * kv_row / 4); + case PosPatch::RowDrain: return uint32_t(pos * kv_row); + case PosPatch::Record: return uint32_t(pos * ptab_row); + case PosPatch::WindowOffset: return 0; + } + throw std::runtime_error("unknown position patch kind"); +} + +void apply_position(std::span w, const std::vector& patches, size_t pos, size_t kv_row, + size_t ptab_row) { + for (const auto& p : patches) { + if (p.word >= w.size()) throw std::runtime_error("position patch past the stream"); + w[p.word] = position_value(p.kind, pos, kv_row, ptab_row) | p.flags; + } +} + +namespace { + +uint32_t le32(const Bytes& b, size_t o) { + if (o + 4 > b.size()) throw std::runtime_error("lax_a ELF: read past the end"); + uint32_t v; + std::memcpy(&v, b.data() + o, 4); + return v; +} + +} // namespace + +std::vector elf_position_sites(const Bytes& e, size_t kv_row) { + const ElfSection ct = elf_section(e, ".ctrltext"), rela = elf_section(e, ".rela.dyn"), + dsym = elf_section(e, ".dynsym"), dstr = elf_section(e, ".dynstr"); + if (ct.size % 4) throw std::runtime_error("lax_a ELF: .ctrltext is not whole words"); + std::vector w(ct.size / 4); + std::memcpy(w.data(), e.data() + ct.offset, ct.size); + + // The BD blockwrite each DDR_PATCH op follows (the last one before it). + std::vector bd_before(w.size(), SIZE_MAX); + for (size_t i = 4, last = SIZE_MAX; i < w.size(); i += op_words(w[i])) { + if (w[i] == 0x81) bd_before[i] = last; + if (w[i] == 0x01) last = i; + } + auto symbol_name = [&](uint32_t sym) { + if (uint64_t(sym + 1) * 16 > dsym.size) throw std::runtime_error("lax_a ELF: relocation symbol out of range"); + const size_t o = dstr.offset + le32(e, dsym.offset + 16u * sym); + const auto* b = reinterpret_cast(e.data()); + return std::string(b + o, strnlen(b + o, e.size() - o)); + }; + + std::vector out; + for (const PosPatch& p : position_patches(w, kv_row)) { + out.push_back({ct.offset + 4 * p.word, p.kind, p.flags, ElfPosSite::CtrlWord, p.word}); + if (p.kind == PosPatch::WindowLength) continue; + // The DDR_PATCH op holding this offset word, and the BD it rewrites. + const size_t op = p.word - 10, bd = bd_before[op]; + if (bd == SIZE_MAX || w[bd + 2] + 4 != w[op + 6]) + throw std::runtime_error("lax_a ELF: no BD write before the patch at word " + std::to_string(p.word)); + const uint32_t bd_offset = uint32_t(4 * (bd + 4)); // the BD's first word, in .ctrltext + const uint32_t arg = w[op + 8]; + int found = 0; + for (size_t k = 0; k < rela.size / 12; ++k) { + const size_t ro = rela.offset + 12 * k; + if (le32(e, ro) != bd_offset) continue; + if (symbol_name(le32(e, ro + 4) >> 8) != std::to_string(arg + 3)) // xclbin numbering + throw std::runtime_error("lax_a ELF: the relocation at word " + std::to_string(bd + 4) + + " names another argument than its DDR_PATCH"); + if ((le32(e, ro + 8) & 0x7FFFFFFFu) != (w[p.word] & 0x7FFFFFFFu)) + throw std::runtime_error("lax_a ELF: addend and DDR_PATCH word disagree at word " + std::to_string(p.word)); + out.push_back({ro + 8, p.kind, le32(e, ro + 8) & 0x80000000u, ElfPosSite::Addend, k}); + ++found; + } + if (found != 1) + throw std::runtime_error("lax_a ELF: " + std::to_string(found) + " relocations on the BD at word " + + std::to_string(bd + 4) + ", expected 1"); + } + return out; +} + +void apply_elf_position(Bytes& e, const std::vector& sites, size_t pos, size_t kv_row, size_t ptab_row) { + for (const auto& s : sites) { + if (s.offset + 4 > e.size()) throw std::runtime_error("position site past the ELF"); + const uint32_t v = position_value(s.kind, pos, kv_row, ptab_row) | s.flags; + std::memcpy(e.data() + s.offset, &v, 4); + } +} + +std::array cfg_words(uint64_t pool_addr) { + std::array c{}; + c[0] = uint32_t(pool_addr); + c[1] = uint32_t(pool_addr >> 32); + for (size_t i = 0; i < kQueues.size(); ++i) c[2 + i] = kQueues[i]; + return c; +} + +} // namespace onebit::npu::lax diff --git a/npu/lax_stream.h b/npu/lax_stream.h new file mode 100644 index 00000000..7880862e --- /dev/null +++ b/npu/lax_stream.h @@ -0,0 +1,103 @@ +// 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. + +// Host-side words the lax decode writes to the device between tokens (docs/npu-lax.md). +// Pure host code, no XRT. +// +// Position patching. The full-attention control text (lax_a insts.bin) is compiled for +// one placeholder position, and all ten full-attention runs of a token share it. An +// mlir-aie instruction stream is 4 header words, then ops: op 1 is a BD blockwrite (12 +// words, its register at +2 and the BD's length at +4), op 0x81 a DDR address patch (12 +// words: register at +6, the kernel's buffer argument at +8, the byte offset into that +// buffer at +10); ops 0, 3 and 0x80 are 6, 7 and 4 words, anything else one. Per token +// four words change: +// the KV window fill's length (the blockwrite before the arg-3 patch at offset 0), in +// 32-bit words: nf rows, nf = max(pos, 1) +// the KV window fill's offset (that patch): the window's first row, always 0 here +// the new row's drain offset (the arg-3 patch at one row): pos rows +// the position record's offset (the arg-5 patch): ptab row pos +// Arguments past the first five carry the firmware's +0x80000000 address translation in +// the offset word; that bit is kept. Layout from the open kernels' harness +// (open_kernels/harness/stream_patch.hpp, attn_table / attn_apply; MIT), re-implemented. +// +// Full ELFs. XRT's ELF flow does not read the DDR_PATCH offset word: aiebu-asm moves +// every offset into the .rela.dyn addend of the relocation on the BD the patch rewrites +// (type 5, at the BD's first word; the BD's address word is zero), and XRT patches the BD +// from bo.address() + that addend. So on the ELF path a position is written at the +// window BD's length word, and at both the DDR_PATCH word and the relocation addend of +// the three offset patches. elf_position_sites derives those seven places from an +// instruction ELF (lax_a's insts.elf) the way position_patches does for insts.bin. +// +// Layer config (cfg, 4 KiB per layer): words 0..1 the layer pool's device address +// (bo.address() + 0x80000000), which the on-device router forms the routed experts' +// addresses from; words 2..9 each column's weight-stream MM2S queue register (the merged +// lax design's shim allocation; open_kernels model/lax_decode_cfg.py QUEUES). +#pragma once + +#include "full_elf.h" + +#include +#include +#include +#include +#include + +namespace onebit::npu::lax { + +struct PosPatch { + enum Kind : uint8_t { WindowLength, RowDrain, Record, WindowOffset }; + size_t word; // index into the instruction words + Kind kind; + uint32_t flags; // bits of the offset word to keep (0x80000000 or 0) +}; + +// Words of one instruction-stream op starting with opcode word w. +size_t op_words(uint32_t w); + +// The four position patches of a full-attention stream; throws unless each is found +// exactly once. kv_row is the KV row in bytes. +std::vector position_patches(std::span words, size_t kv_row); + +// Write position pos into the stream. +void apply_position(std::span words, const std::vector& patches, size_t pos, size_t kv_row, + size_t ptab_row); + +// The value a patch of that kind takes at position pos, before its flags. +uint32_t position_value(PosPatch::Kind kind, size_t pos, size_t kv_row, size_t ptab_row); + +struct ElfPosSite { + enum Where : uint8_t { CtrlWord, Addend }; + size_t offset; // byte offset in the ELF file + PosPatch::Kind kind; + uint32_t flags; // bits to keep (the folded 0x80000000 of args >= 5; full_elf clears it) + Where where; + size_t index; // the .ctrltext word, or the .rela.dyn entry +}; + +// The position sites of a full-attention instruction ELF: the window length word, and +// per offset patch its DDR_PATCH word and the addend of the one relocation on its BD +// (whose symbol must name the patch's argument). Throws unless each is found exactly once. +std::vector elf_position_sites(const Bytes& insts_elf, size_t kv_row); + +// Write position pos into an instruction ELF. +void apply_elf_position(Bytes& insts_elf, const std::vector& sites, size_t pos, + size_t kv_row, size_t ptab_row); + +constexpr std::array kQueues = {0x1D21C, 0x1D214, 0x1D21C, 0x1D21C, 0x1D21C, 0x1D21C, 0x1D214, 0x1D214}; + +// cfg words 0..9 for a pool at device address pool_addr (bo.address() + 0x80000000). +std::array cfg_words(uint64_t pool_addr); + +} // namespace onebit::npu::lax diff --git a/npu/lax_turns.cpp b/npu/lax_turns.cpp new file mode 100644 index 00000000..3e85fea6 --- /dev/null +++ b/npu/lax_turns.cpp @@ -0,0 +1,91 @@ +// 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 "lax_turns.h" + +#include + +namespace onebit::npu::lax { + +namespace { + +bool is_prefix(std::span p, std::span of) { + return p.size() <= of.size() && std::equal(p.begin(), p.end(), of.begin()); +} + +} // namespace + +size_t Snapshots::bytes() const { + size_t n = 0; + for (const auto& e : e_) n += e.state.capacity(); + return n; +} + +int Snapshots::best(std::span prompt) const { + int best = -1; + for (size_t i = 0; i < e_.size(); ++i) { + const auto& t = e_[i].tokens; + if (t.empty() || t.size() >= prompt.size() || !is_prefix(t, prompt)) continue; + if (best < 0 || t.size() > e_[size_t(best)].tokens.size()) best = int(i); + } + return best; +} + +bool Snapshots::contains(std::span tokens) const { + return std::ranges::any_of(e_, [&](const Entry& e) { return std::ranges::equal(e.tokens, tokens); }); +} + +Snapshots::Entry* Snapshots::put(std::span tokens) { + if (cap_ == 0) return nullptr; + Entry* e = nullptr; + for (auto& x : e_) + if (std::ranges::equal(x.tokens, tokens)) e = &x; + if (!e) { + if (e_.size() < cap_) e = &e_.emplace_back(); + else e = &*std::ranges::min_element(e_, {}, &Entry::used); + e->tokens.assign(tokens.begin(), tokens.end()); + } + e->used = ++tick_; + return e; +} + +Plan plan(std::span live, const Snapshots& snaps, std::span prompt) { + Plan p; + if (!live.empty() && live.size() < prompt.size() && is_prefix(live, prompt)) { + p.source = Plan::Live; + p.reuse = live.size(); + } + if (const int i = snaps.best(prompt); i >= 0 && snaps.at(size_t(i)).tokens.size() > p.reuse) { + p.source = Plan::Snapshot; + p.reuse = snaps.at(size_t(i)).tokens.size(); + p.snapshot = i; + } + return p; +} + +std::vector snapshot_points(std::span prompt, size_t from, std::span turn_tokens) { + std::vector out; + if (prompt.size() < 2) return out; + const size_t lo = std::max(from, 1); + for (size_t i = prompt.size() - 1; i-- > 0;) // the last turn token before the last token + if (std::ranges::find(turn_tokens, prompt[i]) != turn_tokens.end()) { + if (i >= lo) out.push_back(i); + break; + } + if (prompt.size() - 1 >= lo) out.push_back(prompt.size() - 1); + return out; +} + +} // namespace onebit::npu::lax diff --git a/npu/lax_turns.h b/npu/lax_turns.h new file mode 100644 index 00000000..31799138 --- /dev/null +++ b/npu/lax_turns.h @@ -0,0 +1,92 @@ +// 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. + +// Where the lax decode's next prompt starts (docs/npu-lax.md, "Chat follow-ups"). Pure +// host bookkeeping, no XRT: lax::Session (npu/lax.h) holds the device side. +// +// A position's cache is the KV rows of the ten full-attention layers and the DeltaNet +// state of the thirty linear ones. The KV rows can be overwritten from any position (a +// position reads rows 0..pos-1 only), but the DeltaNet state folds every token in and +// cannot rewind. So a prompt reuses the device's cache in one of two ways: +// live the prompt extends the tokens on the device: feed the rest; +// snapshot a copy of the DeltaNet state taken at a prefix of the prompt, in host +// memory: load it, set the position to its length, feed the rest; +// and otherwise starts over from position 0. +// +// Snapshots are taken while a prompt is fed, at the points snapshot_points names: before +// its last token (a retry of the same prompt, or one extending it) and before the last +// turn token (<|im_start|>): a chat client's follow-up re-renders the earlier answer +// differently from how it was generated (without its think block), but repeats every +// message before it exactly. +#pragma once + +#include +#include +#include +#include + +namespace onebit::npu::lax { + +// Snapshots of the DeltaNet state at token prefixes, at most `capacity`; the least +// recently used one goes first. +class Snapshots { +public: + struct Entry { + std::vector tokens; // the prefix: positions 0..tokens.size()-1 + std::vector state; // Decoder::save_state's bytes after that prefix + uint64_t used = 0; + }; + + explicit Snapshots(size_t capacity = 0) : cap_(capacity) {} + size_t capacity() const { return cap_; } + size_t size() const { return e_.size(); } + size_t bytes() const; // the host memory the kept states take + const Entry& at(size_t i) const { return e_[i]; } + + // The longest entry whose tokens are a proper prefix of prompt (shorter than it: the + // last prompt token has to run for its logits), or -1. + int best(std::span prompt) const; + bool contains(std::span tokens) const; + void touch(size_t i) { e_[i].used = ++tick_; } + // The entry for tokens, to be filled with its state: the one already there, or a new + // one in the place of the least recently used (whose state buffer it keeps, to avoid + // a reallocation). Its state is left as it was: the caller writes it. nullptr at + // capacity 0. + Entry* put(std::span tokens); + void clear() { e_.clear(); } + +private: + size_t cap_; + uint64_t tick_ = 0; + std::vector e_; +}; + +struct Plan { + enum Source { Scratch, Live, Snapshot }; + Source source = Scratch; + size_t reuse = 0; // prompt tokens not fed again + int snapshot = -1; // with Snapshot: the entry +}; + +// live: the tokens on the device now. The source that feeds the fewest prompt tokens; on a +// tie, live (nothing to load). +Plan plan(std::span live, const Snapshots& snaps, std::span prompt); + +// The prefix lengths at which to take a snapshot while prompt[from..] is fed, ascending: +// prompt.size() - 1, and the index of the last token of the prompt that is one of +// turn_tokens; each only if it is at least max(from, 1). +std::vector snapshot_points(std::span prompt, size_t from, std::span turn_tokens); + +} // namespace onebit::npu::lax diff --git a/npu/tokenizer.h b/npu/tokenizer.h index f733718a..29a42838 100644 --- a/npu/tokenizer.h +++ b/npu/tokenizer.h @@ -41,6 +41,8 @@ class Tokenizer { // Ids outside the vocabulary decode to nothing. std::string decode(const std::vector& ids) const; std::string decode(int id) const { return decode(std::vector{id}); } + // One past the largest id, added tokens included. + int size() const { return int(id_to_text_.size()); } private: std::unordered_map vocab_; // token text -> id, added tokens included diff --git a/scripts/build-lax.sh b/scripts/build-lax.sh new file mode 100755 index 00000000..5e2f370e --- /dev/null +++ b/scripts/build-lax.sh @@ -0,0 +1,102 @@ +#!/usr/bin/env bash +# 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. +# +# build-lax.sh +# +# Builds the Qwen3.6-35B-A3B whole-layer MoE decode (docs/npu-lax.md) from the open +# kernels pinned in third_party/OpenFlowLM-Next, into : +# /kernels/lax_l, lax_a the merged whole-layer design, linear-attention and +# full-attention control texts (one design build) +# /kernels/ln, lm_head_q8 the final RMSNorm and the q8 lm_head +# Each of the four kernel directories holds what both kernel transports need: +# insts.elf, main.pdi full ELFs (the default: npu/lax_elf.h assembles them +# in memory; no xclbin is read) +# final.xclbin, insts.bin the classic path (1bit npu-lax --transport classic) +# and lax_l/lax_a keep their aiecc project in final.prj. +# /src/open_kernels the pinned tree, with its XRT harness built in +# harness/build/run_kernel and the ln / lm_head builds +# where its tools look for them (--designs default) +# The tree is copied to /src first: the kernel generator writes into it, and +# the submodule stays clean. Run the tools from /src/open_kernels (docs/npu-lax.md). +# +# Needs the IRON toolchain the designs are written for (mlir-aie 1.4.2 + Peano, in a +# Python venv: IRON_VENV) and the AIE tools that provide aiebu-asm (AIETOOLS, the +# aietools directory of a Vitis install); xclbinutil on PATH; XRT headers in +# /opt/xilinx/xrt (or XILINX_XRT). Four kernel builds, a few minutes each. +set -euo pipefail +prefix=${1:?usage: build-lax.sh } +: "${IRON_VENV:?set IRON_VENV to the mlir-aie 1.4.2 + Peano venv}" +: "${AIETOOLS:?set AIETOOLS to a Vitis aietools directory (aiebu-asm)}" +root=$(cd "$(dirname "$0")/.." && pwd) +sub=$root/third_party/OpenFlowLM-Next +[ -f "$sub/open_kernels/build_design.py" ] || + { echo "third_party/OpenFlowLM-Next is empty: git submodule update --init third_party/OpenFlowLM-Next"; exit 1; } +mkdir -p "$prefix" +prefix=$(cd "$prefix" && pwd) + +commit=$(git -C "$sub" rev-parse --short=12 HEAD) +rm -rf "$prefix/src" +mkdir -p "$prefix/src" +git -C "$sub" archive HEAD open_kernels | tar -x -C "$prefix/src" +ok=$prefix/src/open_kernels + +# shellcheck disable=SC1091 +source "$IRON_VENV/bin/activate" +export PATH="$AIETOOLS/bin:$VIRTUAL_ENV/bin:/usr/bin:$PATH" +unset PYTHONPATH +spec=recipes/specs/qwen36-35b-a3b.json +cd "$ok" + +# The final norm and the lm_head (the recipe's own build sets). This also regenerates +# the whole-layer designs' kernel sources (designs/layer_x/gen_kernels.py) for the spec. +# --force: /src is fresh on every run, so its build cache key alone proves nothing. +python export_qwen36_kernels.py --spec "$spec" --only ln,lm_head_q8 --out "$prefix/kernels" --force + +# The merged whole-layer design, both control texts from the same source. The three +# ONDV flags are the on-device routing fixes (docs/npu-lax.md, "What made it correct"). +export OPEN_KERNELS_SPEC=$spec MOE_ONDEVICE_ROUTE=1 ONDV_EMIT_SHARED=1 ONDV_PKTDONE_ACQ=1 +LAX_KIND=0 python build_design.py designs/layer_x/lax.py "$prefix/kernels/lax_l" +LAX_KIND=1 python build_design.py designs/layer_x/lax.py "$prefix/kernels/lax_a" +deactivate + +# The tools find the ln / lm_head builds under designs/ (the recipe's manifest), +# which is where the export builds them; check rather than assume. The export copies only +# the classic files (final.xclbin, insts.bin) into kernels/; the full ELFs need each +# build's insts.elf and PDI too, so every kernel directory gets insts.elf and main.pdi. +python3 - "$prefix/kernels" "$ok/designs" <<'PY' +import json, os, shutil, sys +kernels, designs = sys.argv[1], sys.argv[2] +builds = json.load(open(os.path.join(kernels, "manifest.json")))["builds"] +for name in ("ln", "lm_head_q8"): + build = os.path.join(designs, builds[name]["build_dir"]) + for f in ("final.xclbin", "insts.elf", "final.prj/main.pdi"): + if not os.path.isfile(os.path.join(build, f)): + sys.exit(f"{build}/{f} missing after the export") + shutil.copyfile(os.path.join(build, "insts.elf"), os.path.join(kernels, name, "insts.elf")) + shutil.copyfile(os.path.join(build, "final.prj/main.pdi"), os.path.join(kernels, name, "main.pdi")) +for name in ("lax_l", "lax_a"): + shutil.copyfile(os.path.join(kernels, name, "final.prj/main.pdi"), os.path.join(kernels, name, "main.pdi")) +for name in ("lax_l", "lax_a", "ln", "lm_head_q8"): + for f in ("insts.elf", "main.pdi", "final.xclbin", "insts.bin"): + if not os.path.isfile(os.path.join(kernels, name, f)): + sys.exit(f"{kernels}/{name}/{f} missing") +PY + +cmake -S harness -B harness/build -DCMAKE_BUILD_TYPE=Release > /dev/null +cmake --build harness/build -j "$(nproc)" > /dev/null + +echo "lax ($commit): kernels in $prefix/kernels, tools in $ok (harness/build/run_kernel)" +md5sum "$prefix"/kernels/lax_l/insts.bin "$prefix"/kernels/lax_a/insts.bin "$prefix"/kernels/lax_a/insts.elf diff --git a/tests/golden/npu_lax/lax_a_elf_sites.tsv b/tests/golden/npu_lax/lax_a_elf_sites.tsv new file mode 100644 index 00000000..e2cad6d3 --- /dev/null +++ b/tests/golden/npu_lax/lax_a_elf_sites.tsv @@ -0,0 +1,10 @@ +# the full-ELF position sites of lax_a insts.elf md5 426b8dd11adf760697b1c1815fe64f0e (scripts/build-lax.sh at the pin), as the +# lax-elf investigation's attn_sites() found them (OpenFlowLM-Next 1bit/lax-elf 65ad44d, harness/full_elf.py) +# where index (.ctrltext word | .rela.dyn entry) kind (0 window length, 1 row drain, 2 record, 3 window offset) flags +word 2762 1 00000000 +reloc 68 1 00000000 +word 2984 2 80000000 +reloc 74 2 80000000 +word 3222 0 00000000 +word 3240 3 00000000 +reloc 78 3 00000000 diff --git a/tests/golden/npu_lax/lax_a_patches.tsv b/tests/golden/npu_lax/lax_a_patches.tsv new file mode 100644 index 00000000..a3b7607d --- /dev/null +++ b/tests/golden/npu_lax/lax_a_patches.tsv @@ -0,0 +1,6 @@ +# the reference harness's attn_table (open_kernels/harness/stream_patch.hpp, kv_row 2048) for lax_a insts.bin md5 4f3c749d8b0263260b8482c7350e7ebc (scripts/build-lax.sh at the pin) +# word kind (0 window length, 1 row drain, 2 record, 3 window offset) flags +2762 1 00000000 +2984 2 80000000 +3222 0 00000000 +3240 3 00000000 diff --git a/tests/golden/npu_lax/sha256.tsv b/tests/golden/npu_lax/sha256.tsv new file mode 100644 index 00000000..1bc95a8f --- /dev/null +++ b/tests/golden/npu_lax/sha256.tsv @@ -0,0 +1,83 @@ +ptab 4194304 2a11c3c044d7cd4cd3a07abb7c2b5e40e6fe11b2793c891b44ccf5e28d2ad8f7 +normw 4096 43a07200ff70a20e7901e8129095e06ae89936cc772ee76fc5f2a3aa4557260c +lmpool 542113792 84c1d04930e651a48b7172760484e3b5f397201ecef96929f6e51464eac0f4c2 +pool_0 536870912 5d31965a106d50cd6d84965c77fd9121aa096ae367edc510e65f0be995a50d83 +consts_0 11882496 3d3b24fd353638d2ccdf611fdc0bdb244b20a180b282208fd50e43c584d6d808 +pool_1 536870912 ccbd30d76159416d3731dbbf973f6a3ad7a93d1438a1b8ff478f0b727f8f0ca7 +consts_1 11882496 2a1e96820ceed2d96d33a076314391eb627236453aaad9344445980769a21408 +pool_2 536870912 3aeda643c0285dc2ab0ce3d0c17ae20d024213bb23a163f032aded77ff3c46ce +consts_2 11882496 c6be2bf18907f8c2659d2cd3d906d89f36b5fbb1adfd7ec6aa4cb01e1dc58d53 +pool_3 536870912 f0aac475e7fcf048ed10ba8886349a32f8a5302d86e5cbdff4453a73e04d619f +consts_3 1062912 bfd3d20bb30300898c3bd2774ef6004ef87c8ebe9d96dec3db3baea236748525 +pool_4 536870912 bb4987d12ce6101764e91be7774c96fe7f215a58040226297e559e0443f9860f +consts_4 11882496 087150c1656c8268e7dd8ebd24804ef75797f8e260738435526847dc444cd669 +pool_5 536870912 19de3ad9652f5d6d2276b79dc073ab2da53ee780e6daaeefd6dd55416154ee1f +consts_5 11882496 7dc6c146a0a146b6ac2072fb322252aeed07d68f1d6d136c1383233dd9c9a492 +pool_6 536870912 3966fdeb81fd87f1288b4fd1aa5baaab1cffd21a7de99e4b485a82eeb1ac62a3 +consts_6 11882496 0c0855742b266f1d2f7cf61734cbc3a1153d59d45b4c8839d1b40e45bf139a5a +pool_7 536870912 cfee56e8b9b1040398354491d1216f2ff8e149e2b5451d01da9b29664077e773 +consts_7 1062912 00aab7349d0fa8958003af1e406fe6e029412172596d749afa5860b51045baa1 +pool_8 536870912 832676f42a75f6088021fc22c1c9b089d0dba97e64c31a101e74116044d906d0 +consts_8 11882496 f69e3b5e738998f27c17ab39448392e41cee9d25fd3f2bf1e969d125d0ecc09e +pool_9 536870912 8b30b72f5ca34af0a32f11ad2f40604db1315f40f1152a3f0cd3a46ec445eaa4 +consts_9 11882496 7f7838ecdb8f641723740bad93e2d704d542fbc95db1b5a4d1285e9668caa35c +pool_10 536870912 634dd8c8ed885e63585eb77a1cf2914cae7706d652735f1f1b640329ff357415 +consts_10 11882496 3a4739daee15c758e65915836bd7c354c01765c64b145cff6f8c6b1325826566 +pool_11 536870912 414aea8c275cbce6d046561b337e3bb80b6d2d6783c6a35ed7b12ea327a6f45f +consts_11 1062912 9b31ba82ce610e03e4fcaba423e60270e721be38d85a2f66ff65f1d33cd2948e +pool_12 536870912 571b1da81ae63b3b39ae7bee1d081829ead789b3bf398080febc89eaca412ffe +consts_12 11882496 8c90941fe2b3098f134ca3f70f7b559fd8a053e670be0c70889195c2f5ee6a12 +pool_13 536870912 34f3697fbedbc0d15f28da0b326da245ae90ebbac4662b11fae882c12765072e +consts_13 11882496 65204617994cd61b2d748ce66ec2d96b56e3ecc6e2bb42f0bd80d76a0a7b9887 +pool_14 536870912 31ed19a3d6fc665d8ba27065c0b5fb0aecf0d4f03ae83213d80f37ecc5668ddc +consts_14 11882496 2b998f4ca3e2534edb68dbf75c74d2fc4f9eea7044fd0c7d5053be01d0290b0c +pool_15 536870912 abc631cb97f5cb50d8c73c68094352e0ca9889a39414fe4d7805afca53a1264b +consts_15 1062912 c767e01505333aee11e29593603d2ba743d2af706e13e00fd46ba344fde5b6da +pool_16 536870912 4adae192a12395532220f04a9df594d7c1d57d06035efac17022253008f4aa31 +consts_16 11882496 80a194d49424b310706876c8033f59c4937c9fd449c55b80b9a3ab79f4f3bb24 +pool_17 536870912 853a93791f1895dc53b51fde8ec39256612a60bf298aa70620ee494ab64e7a64 +consts_17 11882496 fef059c14e28690992a3feb2925af3d66919c892851d0c5e1d2c28dacab63267 +pool_18 536870912 237f4f77352b893de39afc786e31efac3b840db4bef1b50673f4163242242e00 +consts_18 11882496 db58de73a6cf040d7d150359d06019301f10088c5869d18fb5cf8ea318e4b2c9 +pool_19 536870912 5744885732eebdd8681f6147200a754f3c566752dffbf6f3ab82f7c2eaa88098 +consts_19 1062912 4635a58f00859440f434e468d67785b573b470f4125ab59cf4c220d3800ed4f1 +pool_20 536870912 94321430e998460606b1f92290539039022e3c4fdd2359be93812db82b257803 +consts_20 11882496 cc5c1209e7a9ea714f58abee915119624ea24b246fece8a1a7248e01ff9b9099 +pool_21 536870912 b28d9325abf06c8d81904955299c8b2b04a34a8d51aff3e08c7b2590a58592cd +consts_21 11882496 2ab2f5469cabc94eb3fadb385da416b3ef675ee4e7db64d89b11ef5c78b58531 +pool_22 536870912 2eef0ee382d24ad46d8376ae7e5bef618b0320fea32d61e591096d82237cf629 +consts_22 11882496 66590b7c59fa92b86295025f01a62b9dbb74462da5712a4e4476bf4d17171d8c +pool_23 536870912 7d5c708586aab7833f0303378c7ecfa17ac91f4b5db0f4c6e08a6f05996ef8be +consts_23 1062912 1d29bd66faaa23a196ecc69aa4506241cd7f6c7aa405765339748147c7cae4f1 +pool_24 536870912 e95c5a5e3df9fcca7affe5530c863a934449002a0b4f64fdfa44d4e347dea700 +consts_24 11882496 1ac63c2ea432661ffb1968c9d065d97274346e2ecb6464bb8bfd6fcf9b260c00 +pool_25 536870912 a8ad3ecb18f60596a41de8dc954444f4d8eb01e93f0999a1f1ac8dfde64c0185 +consts_25 11882496 5acc5e073914ee19642ed9117306cb9e706712a05f6b947fec2291fd8cebd4f5 +pool_26 536870912 205bb69efda97f3d2cdb7b91ab2103f80296dcc60d3f25e9479ccf3f4b707d5f +consts_26 11882496 f65ec624f5ee09f8bdba34245a78fe261c861bb09940d3e9b2909de9e34f9007 +pool_27 536870912 c2514b518a0657bae083ea26f0370305f0c5789fe1b77e7142b395abe529a294 +consts_27 1062912 3ea810440bd0c22f2dc66ddb817e8baecf356fa0f9353ee113ca2e05bcc8f7d9 +pool_28 536870912 4242535074a8baae40771f94efdeffbc985a058b0fdfe82a95e6db7f4e04dd48 +consts_28 11882496 601888690f15d0be5cc3caa7e147c6d2cd4cba6e0406219997bf6b46beea0157 +pool_29 536870912 40f7495f1ede805991b5cfddbb83aced3dc5d2907e0d6bc54bfa3c7c5c39aefa +consts_29 11882496 379f1a4da443b75e084272c490a6bc9b6914cb0c61bcb47499c88a8bd670bee7 +pool_30 536870912 e796b21887a85068817a28df5a21bc2f1a7a21feeb38458240be8e8a275f020f +consts_30 11882496 ae9f05766e7a35d7bfb972b5063ad29cdc37686d2917c6059023c910eed9da4f +pool_31 536870912 d28e956df63dc6a402e83bd23b08708f4ffd2183a5068ad8c3bf9023bae78bcd +consts_31 1062912 06ccb86382b83e9d8f60d04c39b5864647913ac724719bdd14f71fdc6a4e8788 +pool_32 536870912 8d72a4b9cdbf40055cdd420bc179e25e722de3b79cd9f5c53cb750cbc4916c87 +consts_32 11882496 d82e763b5b07f42908664750ed8d538efab6ba9a8dbdc2d2f71ab51f45f12d09 +pool_33 536870912 51e6f52e5e82cbfe52e39e445179c76bcc6e883e728ec2141831bec4767b960f +consts_33 11882496 dfcbc4b51c6078cc634c9a28f65dd20d03767129c10a2bb67502f208f769ec99 +pool_34 536870912 adc7e3177207098e9d4cde688f4171fd9bf2674075d0e60d7a7089e0980f1fe9 +consts_34 11882496 db1cd5360798d66b1de4788a93ae564e0c8f87167f61ef72b4c7b7654351b537 +pool_35 536870912 27c452dfa5230cefbb052d824f664fdb3ae91e090fbed11454867ac1212e2c66 +consts_35 1062912 afa95c240c9b4c4f72a0c6f9a57474e1bbc885c25ecc521e2b786c7d05d0b8c0 +pool_36 536870912 1e817c5b1e91281ebaa6022cdb51fd29eede329233f3001e24dca994b3d2340c +consts_36 11882496 3a5c34e196bbd82445bf947b24cabf1935abd770a33c895d0c0b06b8db7d7acd +pool_37 536870912 d6d899dcd30d73648be5b9576a238c59aa3d801e15f35209e76ab3e30a887b8a +consts_37 11882496 7873ec9ec588433ad0e7674fac53815ff272276dc09742abe8197d9c7968c941 +pool_38 536870912 8f7375f5676728d33386745f88370540e82b82bfb1b2a9d1adad2c2bc4d1656d +consts_38 11882496 95b01ec8d8814e876d4c0bd835c9d005061b30ac6e0cc78ad1f2d4d780f181ee +pool_39 536870912 eb4cacaed9d50be966440bd63b5cabd21d9a529e144457a78b537ea674cf7695 +consts_39 1062912 a44dd665a72b82c1d263eab912067ad62408eaece73cdbc46c3894bceb198146 diff --git a/tests/npu_lax_cpp.sh b/tests/npu_lax_cpp.sh new file mode 100755 index 00000000..29f2a477 --- /dev/null +++ b/tests/npu_lax_cpp.sh @@ -0,0 +1,36 @@ +#!/usr/bin/env bash +# 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. +# +# npu_lax_cpp.sh <1bit> [elf|classic] +# +# The Qwen3.6-35B-A3B lax decode driven by the engine itself (`1bit npu-lax`, no Python): +# three positions from the reference's residual inputs, each one runlist submit plus the +# norm and the lm head, scored against the fp64 reference (corr > 0.9999, same argmax); +# then again after a reset (bit-identical logits); then one greedy chat turn, which must +# answer with Paris. The kernel transport is full ELFs unless classic is given. +# scripts/build-lax.sh's /kernels +# open_kernels/model/make_decode.py --requant --tokens 3 output +# The NPU is shared: run it under the box's lock, e.g. flock ctest -R npu_lax_e2e. +set -euo pipefail +bin=${1:?usage: npu_lax_cpp.sh <1bit> } +model=${2:?model dir} kernels=${3:?kernel dir} ref=${4:?reference dir} transport=${5:-elf} + +"$bin" npu-lax --model "$model" --kernels "$kernels" --transport "$transport" --parity "$ref" --tokens 3 + +answer=$("$bin" npu-lax --model "$model" --kernels "$kernels" --transport "$transport" -n 32 \ + "What is the capital of France? Answer in one sentence.") +echo "chat: $answer" +grep -q Paris <<< "$answer" diff --git a/tests/npu_lax_parity.sh b/tests/npu_lax_parity.sh new file mode 100755 index 00000000..994dc3df --- /dev/null +++ b/tests/npu_lax_parity.sh @@ -0,0 +1,44 @@ +#!/usr/bin/env bash +# 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. +# +# npu_lax_parity.sh +# +# The Qwen3.6-35B-A3B lax decode (docs/npu-lax.md) against the fp64 reference: three +# positions from <|im_start|>, all 40 layers as ONE runlist submit per token, weights +# packed at load time from /model.q4nx. Every position must have finite +# logits, correlation > 0.9999 with the reference and the reference's argmax. Then one +# chat turn must mention Paris. +# scripts/build-lax.sh's prefix +# open_kernels/model/make_decode.py --requant --layers 40 --tokens 3 output +set -euo pipefail +prefix=$(cd "${1:?usage: npu_lax_parity.sh }" && pwd) +model=$(cd "${2:?model dir}" && pwd) +ref=$(cd "${3:?reference dir}" && pwd) +ok=$prefix/src/open_kernels +out=$(mktemp -d) +trap 'rm -rf "$out"' EXIT + +cd "$ok" +python3 model/lax_decode_cfg.py --model-dir "$model" --ref "$ref" --out "$out" \ + --lax-l "$prefix/kernels/lax_l" --lax-a "$prefix/kernels/lax_a" --per 40 --tokens 3 +ln -s "$ref"/ref_logits*.bin "$out"/ +python3 model/compare_decode.py --out "$out" --tokens 3 + +answer=$(python3 model/lax_chat.py --model-dir "$model" --lax-l "$prefix/kernels/lax_l" \ + --lax-a "$prefix/kernels/lax_a" --max-new 32 \ + "What is the capital of France? Answer in one sentence." 2>/dev/null) +echo "chat: $answer" +grep -q Paris <<< "$answer" diff --git a/tests/npu_lax_serve.sh b/tests/npu_lax_serve.sh new file mode 100755 index 00000000..9a62dfbf --- /dev/null +++ b/tests/npu_lax_serve.sh @@ -0,0 +1,134 @@ +#!/usr/bin/env bash +# 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. +# +# `1bit serve --device npu` on Qwen3.6-35B-A3B (docs/npu-lax.md, docs/serve.md): the +# model directory routes to the lax decode on full ELFs, in process. Under strace, it +# checks /health, /v1/models, a chat that answers "The capital of France is Paris.", the +# same request again (it restores the DeltaNet state snapshot taken before the prompt's last +# token and must answer the same), the chat follow-up a client sends (the history +# re-rendered: it restores the snapshot taken before the first prompt's assistant header), +# a raw prompt that extends the last one (the live cache is reused and it answers Berlin), a +# streamed chat with its usage chunk, and that no .xclbin was opened while the lax kernels' +# insts.elf were. +# +# usage: tests/npu_lax_serve.sh <1bit> [extra serve args] +# The NPU is shared: run it under the box's lock, e.g. flock ctest -R npu_lax_serve. +set -uo pipefail + +bin=${1:?usage: npu_lax_serve.sh <1bit> [args]} +model=${2:?model dir} kernels=${3:?kernel dir} +shift 3 +port=$(python3 -c 'import socket; s=socket.socket(); s.bind(("127.0.0.1", 0)); print(s.getsockname()[1])') +log=$(mktemp) trace=$(mktemp) +tracer=() +if command -v strace >/dev/null; then tracer=(strace -f -qq --seccomp-bpf -e trace=open,openat -o "$trace"); fi +"${tracer[@]}" "$bin" serve -m "$model" --device npu --npu-kernels "$kernels" --port "$port" --alias lax-e2e "$@" \ + >"$log" 2>&1 & +pid=$! +stop() { pkill -TERM -P "$pid" 2>/dev/null; kill "$pid" 2>/dev/null; wait "$pid" 2>/dev/null; } # 1bit, then strace +cleanup() { stop; rm -f "$log" "$trace"; } +trap cleanup EXIT +fail=0 +check() { if eval "$2"; then echo "ok $1"; else echo "FAIL $1"; fail=1; fi; } +api="http://127.0.0.1:$port" + +code=000 +t0=$(date +%s.%N) +for _ in $(seq 1 600); do + code=$(curl -s -o /dev/null -w '%{http_code}' "$api/health") + [ "$code" = 200 ] && break + kill -0 "$pid" 2>/dev/null || break + sleep 0.2 +done +echo " ready after $(python3 -c "import sys; print(f'{float(sys.argv[2]) - float(sys.argv[1]):.1f}')" "$t0" "$(date +%s.%N)") s" +check "/health 200" '[ "$code" = 200 ]' +models=$(curl -s "$api/v1/models") +check "/v1/models names lax-e2e" '[[ "$models" == *lax-e2e* ]]' + +body='{"model": "lax-e2e", "messages": [{"role": "user", "content": "What is the capital of France? Answer in one sentence."}], + "max_tokens": 32, "chat_template_kwargs": {"enable_thinking": false}}' +ask() { + local reply + reply=$(curl -s "$api/v1/chat/completions" -H 'Content-Type: application/json' -d "$body") + python3 -c ' +import json, sys +r = json.loads(sys.argv[1]) +t, u = r.get("timings", {}), r["usage"] +print(r["choices"][0]["message"]["content"]) +print("%d prompt tokens in %.0f ms, %d generated in %.0f ms" % (u["prompt_tokens"], t.get("prompt_ms", 0), + u["completion_tokens"], t.get("predicted_ms", 0))) +print(u["prompt_tokens"], u["prompt_tokens_details"]["cached_tokens"]) +' "$reply" 2>/dev/null || echo "" +} +reply=$(ask) +first=$(head -1 <<< "$reply") +read -r n1 _ <<< "$(tail -1 <<< "$reply")" +echo " said: $first ($(sed -n 2p <<< "$reply"))" +check "answers \"The capital of France is Paris.\"" '[ "$first" = "The capital of France is Paris." ]' +reply=$(ask) +read -r n2 c2 <<< "$(tail -1 <<< "$reply")" +echo " again: $(head -1 <<< "$reply") ($(sed -n 2p <<< "$reply"), $c2 cached)" +check "the same request again answers the same" '[ "$(head -1 <<< "$reply")" = "$first" ]' +check "and restores the snapshot before its last token ($c2 of $n2 cached)" '[ "${c2:-x}" = "$((n2 - 1))" ]' +field() { python3 -c 'import json,sys; r=json.loads(sys.argv[1]); print(eval(sys.argv[2]))' "$1" "$2" 2>/dev/null; } +# The follow-up a chat client sends (the history re-rendered) after that turn: it cannot +# extend the device's tokens (the answer comes back without the think block it was +# generated after), but the first turn's user message is a snapshot's prefix. +r3=$(curl -s "$api/v1/chat/completions" -H 'Content-Type: application/json' -d '{"messages": [ + {"role": "user", "content": "What is the capital of France? Answer in one sentence."}, + {"role": "assistant", "content": "The capital of France is Paris."}, + {"role": "user", "content": "And of Germany?"}], "max_tokens": 32, "chat_template_kwargs": {"enable_thinking": false}}') +c3=$(field "$r3" 'r["usage"]["prompt_tokens_details"]["cached_tokens"]') +echo " chat follow-up: $(field "$r3" 'r["choices"][0]["message"]["content"]') ($c3 of $(field "$r3" 'r["usage"]["prompt_tokens"]') prompt tokens cached, from $(field "$r3" 'r["timings"].get("cache")'))" +# The first prompt's user message: all of it but "<|im_start|>assistant\n\n\n\n\n". +check "the chat follow-up reuses the first turn's user message ($c3 cached)" '[ "${c3:-0}" -ge $((n1 - 8)) ] && [ "${c3:-0}" -lt "$n1" ]' +check "and answers Berlin" '[[ "$r3" == *Berlin* ]]' +# A raw prompt that extends the last one's tokens continues the device's conversation. +p1=$'<|im_start|>user\nWhat is the capital of France? Answer in one sentence.<|im_end|>\n<|im_start|>assistant\n\n\n\n\n' +complete() { + curl -s "$api/v1/completions" -H 'Content-Type: application/json' \ + -d "$(python3 -c 'import json,sys; print(json.dumps({"prompt": sys.argv[1], "max_tokens": 32}))' "$1")" +} +r1=$(complete "$p1") +a1=$(field "$r1" 'r["choices"][0]["text"]') +p2="$p1$a1"$'<|im_end|>\n<|im_start|>user\nAnd of Germany?<|im_end|>\n<|im_start|>assistant\n\n\n\n\n' +r2=$(complete "$p2") +cached=$(field "$r2" 'r["usage"]["prompt_tokens_details"]["cached_tokens"]') +echo " continued: $(field "$r2" 'r["choices"][0]["text"]') ($cached of $(field "$r2" 'r["usage"]["prompt_tokens"]') prompt tokens cached)" +want=$(field "$r1" 'r["usage"]["prompt_tokens"] + r["usage"]["completion_tokens"]') +check "a prompt extending the last one reuses the device's cache (all $want tokens of the last turn)" '[ "${cached:-x}" = "$want" ]' +check "and answers Berlin" '[[ "$r2" == *Berlin* ]]' +sse=$(curl -sN "$api/v1/chat/completions" -H 'Content-Type: application/json' \ + -d '{"model": "lax-e2e", "stream": true, "stream_options": {"include_usage": true}, + "messages": [{"role": "user", "content": "Count to five."}], + "max_tokens": 16, "chat_template_kwargs": {"enable_thinking": false}}') +chunks=$(grep -c '^data: {' <<< "$sse") +check "streams ($chunks chunks)" '[ "$chunks" -gt 3 ]' +check "with a usage chunk last (stream_options.include_usage)" \ + '[[ "$(grep "^data: {" <<< "$sse" | tail -1)" == *'"'"'"prompt_tokens"'"'"'* ]]' + +stop +if [ ${#tracer[@]} -gt 0 ]; then + xclbins=$(grep -c '\.xclbin' "$trace") + elfs=$(grep -c 'insts\.elf"' "$trace") + check "no .xclbin opened ($xclbins), lax insts.elf opened ($elfs)" '[ "$xclbins" = 0 ] && [ "$elfs" -ge 4 ]' +else + echo "skip strace not installed: the no-xclbin check did not run" +fi + +grep -E "lax decode|ready in" "$log" | sed 's/^/ /' +if [ $fail -ne 0 ]; then echo "--- serve log (tail)"; tail -30 "$log"; echo FAIL; exit 1; fi +echo PASS diff --git a/tests/npu_lax_test.cpp b/tests/npu_lax_test.cpp new file mode 100644 index 00000000..e763a60a --- /dev/null +++ b/tests/npu_lax_test.cpp @@ -0,0 +1,722 @@ +// 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. + +// npu_lax_test (CI: pure host, no model) +// npu_lax_test --model --sha256 [--ref ] [--only a,b,...] [--out ] [--threads N] +// npu_lax_test --insts --patches +// npu_lax_test --elf --sites [--golden ] +// +// The lax decode's host side (npu/lax_pack.h, npu/lax_stream.h) against the reference +// packer, the open kernels' Python (third_party/OpenFlowLM-Next open_kernels/recipes/ +// pack.py at the pin; tests/golden/npu_lax/ holds what it produced): +// +// CI: SHA-256 test vectors; the q8 -> q4_1 re-quantization and the signed-nibble +// transcode on synthetic chunks from a 64-bit LCG (x = x * 6364136223846793005 + +// 1442695040888963407, byte = x >> 56), hashed against the reference's output for the +// same bytes; every chunk permutation's index array (int64) hashed against the +// reference's; the 4096-row position table; the position patches on a synthetic stream; +// the cfg words. +// +// --model: packs every buffer of the model (ptab, normw, lmpool, pool_, consts_) +// and compares its SHA-256 with the reference's table; with --ref, the files there +// (pool_0.bin, ...) are also compared byte for byte, naming the first difference; --out +// writes ours there. +// +// --insts: the position patches found in the lax_a stream equal the tsv (word, kind, +// flags) the reference harness's table produced for the same build. +// +// CI also runs the full-ELF position sites (npu/lax_stream.h, elf_position_sites) on a +// synthetic instruction ELF: which words and addends move, their values at a position, +// the flag bit kept until assemble_full_elf clears it, and the malformed inputs refused. +// +// CI also runs the chat follow-up bookkeeping (npu/lax_turns.h): the snapshot store's +// longest-proper-prefix match, LRU eviction and buffer reuse, the live-vs-snapshot plan, +// the snapshot points, and a three-turn chat replayed on a stand-in decoder whose state is +// a hash of the tokens fed: every restored snapshot must equal that prefix fed from +// scratch, and each follow-up must reuse the whole history up to its earlier answer. +// +// --elf: the position sites of the kernel directory's lax_a/insts.elf equal the tsv (the +// investigation's attn_sites for the same build); with --golden, every full ELF +// (npu/lax_elf.h) is byte-identical to the one the reference harness assembled and ran +// (harness/full_elf.py lax: lax_init, lxf, axf_p, ln, lm). +#include "full_elf.h" +#include "lax_elf.h" +#include "lax_pack.h" +#include "lax_stream.h" +#include "lax_turns.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +using namespace onebit::npu; +using namespace onebit::npu::lax; + +namespace { + +int failures = 0; +std::mutex mu; +void check(bool ok, const std::string& what) { + std::lock_guard lk(mu); + if (!ok) ++failures; + std::printf("%s %s\n", ok ? "ok " : "FAIL", what.c_str()); +} + +std::string sha(const std::vector& v) { return hex(sha256(v)); } +std::string sha_str(const std::string& s) { return hex(sha256({reinterpret_cast(s.data()), s.size()})); } + +struct Lcg { + uint64_t x; + uint8_t byte() { + x = x * 6364136223846793005ull + 1442695040888963407ull; + return uint8_t(x >> 56); + } +}; + +std::vector q8_chunks(int n) { + Lcg g{1}; + std::vector out(size_t(n) * kQ8, 0); + for (int c = 0; c < n; ++c) { + uint8_t* p = out.data() + size_t(c) * kQ8; + for (int k = 0; k < 256; ++k) { + uint16_t u = uint16_t((c % 4 != 3 ? 0x3A00 : 0xBA00) + g.byte()); + if (c == 4 && k % 3 == 0) u = 0; // zero-scale blocks: +0 and -0 values, the min's sign + p[2 * k] = uint8_t(u), p[2 * k + 1] = uint8_t(u >> 8); + } + for (int i = 0; i < 8192; ++i) { + uint8_t b = g.byte(); + if (c == 1) b = 0; // constant blocks: d = 0 + if (c == 2) b &= 0x7F; // non-negative codes + p[512 + i] = b; + } + } + return out; +} + +std::vector q4_chunks(int n) { + Lcg g{2}; + std::vector out(size_t(n) * kQ41, 0); + for (int c = 0; c < n; ++c) { + uint8_t* p = out.data() + size_t(c) * kQ41; + for (int k = 0; k < 256; ++k) { + const uint16_t u = uint16_t(0x3A00 + g.byte()); + p[2 * k] = uint8_t(u), p[2 * k + 1] = uint8_t(u >> 8); + } + for (size_t i = 1024; i < kQ41; ++i) p[i] = g.byte(); + } + return out; +} + +template std::string perm_sha(size_t n, F f) { + std::vector b(n * 8); + for (size_t i = 0; i < n; ++i) { + const int64_t v = int64_t(f(i)); + std::memcpy(b.data() + 8 * i, &v, 8); + } + return sha(b); +} + +// ---- a synthetic instruction ELF (the layout aiebu-asm writes) ---------------------------- + +struct Rel { + uint32_t offset, sym; + uint32_t addend; +}; + +// ELF32 LE with .shstrtab, .ctrltext, .dynstr, .dynsym (symbol 0 null, then one per +// argument name) and .rela.dyn: what elf_position_sites and assemble_full_elf read. +Bytes make_insts_elf(const std::vector& ctrl, const std::vector& args, const std::vector& rels) { + auto put32 = [](Bytes& b, uint32_t v) { for (int i = 0; i < 4; ++i) b.push_back(uint8_t(v >> (8 * i))); }; + auto put16 = [](Bytes& b, uint16_t v) { b.push_back(uint8_t(v)); b.push_back(uint8_t(v >> 8)); }; + Bytes shstr = {0}, ct, dynstr = {0}, dynsym(16, 0), rela; + const std::vector names = {"", ".shstrtab", ".ctrltext", ".dynstr", ".dynsym", ".rela.dyn"}; + std::vector name_off; + for (const auto& n : names) { + name_off.push_back(uint32_t(shstr.size())); + shstr.insert(shstr.end(), n.begin(), n.end()); + shstr.push_back(0); + } + for (uint32_t v : ctrl) put32(ct, v); + for (const auto& a : args) { + put32(dynsym, uint32_t(dynstr.size())); + put32(dynsym, 0), put32(dynsym, 0); + dynsym.push_back(0x11), dynsym.push_back(0), put16(dynsym, 0); + dynstr.insert(dynstr.end(), a.begin(), a.end()); + dynstr.push_back(0); + } + for (const auto& r : rels) put32(rela, r.offset), put32(rela, (r.sym << 8) | 5), put32(rela, r.addend); + Bytes e(52, 0); + const std::vector blobs = {&shstr, &ct, &dynstr, &dynsym, &rela}; + std::vector off; + for (const Bytes* b : blobs) { + while (e.size() % 4) e.push_back(0); + off.push_back(uint32_t(e.size())); + e.insert(e.end(), b->begin(), b->end()); + } + while (e.size() % 4) e.push_back(0); + const uint32_t shoff = uint32_t(e.size()); + e.resize(e.size() + 40 * names.size(), 0); + const uint8_t ident[] = {0x7f, 'E', 'L', 'F', 1, 1, 1}; + std::memcpy(e.data(), ident, sizeof ident); + auto w32 = [&](size_t o, uint32_t v) { std::memcpy(e.data() + o, &v, 4); }; + auto w16 = [&](size_t o, uint16_t v) { std::memcpy(e.data() + o, &v, 2); }; + w32(0x20, shoff); + w16(0x2E, 40), w16(0x30, uint16_t(names.size())), w16(0x32, 1); + const uint32_t types[] = {0, 3, 1, 3, 11, 4}; + for (size_t i = 1; i < names.size(); ++i) { + const size_t h = shoff + 40 * i; + w32(h, name_off[i]), w32(h + 4, types[i]), w32(h + 16, off[i - 1]), w32(h + 20, uint32_t(blobs[i - 1]->size())); + } + return e; +} + +uint32_t rd(const Bytes& b, size_t o) { + uint32_t v; + std::memcpy(&v, b.data() + o, 4); + return v; +} + +void elf_sites_unit() { + // TXN header, then per transfer the BD blockwrite (register at +2, BD words from +4: + // length, address = 0) and the DDR_PATCH on the BD's address register (arg at +8, + // offset at +10). Transfers: the KV window fill (arg 3, offset 0), a weight fill (arg + // 0), the row drain (arg 3, one row), the record fill (arg 5, folded: bit 31 set). + std::vector w = {0x06040100, 0, 0, 0}; + std::vector rels; + // symbols: 1 -> "3" (arg 0), 2 -> "6" (arg 3), 3 -> "8" (arg 5) + auto xfer = [&](uint32_t reg, uint32_t sym, uint32_t arg, uint32_t len, uint32_t off) { + const size_t bd = w.size(); + w.insert(w.end(), {0x01, 0, reg, 32, len, 0, 0, 0, 0, 0, 0, 0}); + w.insert(w.end(), {0x81, 0, 0, 0, 0, 0, reg + 4, 0, arg, 0, off, 0}); + rels.push_back({uint32_t(4 * (bd + 4)), sym, off}); + }; + xfer(0x1D000, 2, 3, 512, 0); + w.insert(w.end(), {0x00, 0, 0, 0, 0, 0}); + xfer(0x1D020, 1, 0, 99, 0x1E1E0000); + xfer(0x1D040, 2, 3, 64, uint32_t(kKvRow)); + w.insert(w.end(), {0x80, 0, 0, 0}); + xfer(0x1D060, 3, 5, 64, uint32_t(kPtabRow) | 0x80000000u); + w[2] = 10, w[3] = uint32_t(4 * w.size()); + const std::vector args = {"3", "6", "8"}; + Bytes e = make_insts_elf(w, args, rels); + + const auto sites = elf_position_sites(e, kKvRow); + check(sites.size() == 7, "ELF: 7 position sites (window length, 3 x (DDR_PATCH word, addend))"); + auto find = [&](PosPatch::Kind k, ElfPosSite::Where where) { + for (const auto& x : sites) + if (x.kind == k && x.where == where) return x; + throw std::runtime_error("site missing"); + }; + check(find(PosPatch::WindowLength, ElfPosSite::CtrlWord).index == 8, "ELF: window length = the fill BD's first word"); + check(find(PosPatch::WindowOffset, ElfPosSite::CtrlWord).index == 26 && + find(PosPatch::WindowOffset, ElfPosSite::Addend).index == 0, + "ELF: window offset = DDR_PATCH word 26 + relocation 0"); + check(find(PosPatch::RowDrain, ElfPosSite::CtrlWord).index == 80 && + find(PosPatch::RowDrain, ElfPosSite::Addend).index == 2, + "ELF: row drain = its DDR_PATCH word + relocation 2 (the weight fill's is skipped)"); + const auto rec = find(PosPatch::Record, ElfPosSite::Addend); + check(rec.index == 3 && rec.flags == 0x80000000u, "ELF: record = relocation 3, folded bit kept"); + + apply_elf_position(e, sites, 7, kKvRow, kPtabRow); + const ElfSection ct = elf_section(e, ".ctrltext"), rl = elf_section(e, ".rela.dyn"); + auto word = [&](size_t i) { return rd(e, ct.offset + 4 * i); }; + auto addend = [&](size_t k) { return rd(e, rl.offset + 12 * k + 8); }; + check(word(8) == 7 * kKvRow / 4, "ELF position 7: window length 7 rows in words"); + check(word(26) == 0 && addend(0) == 0, "ELF position 7: window from row 0 (word and addend)"); + check(word(80) == 7 * kKvRow && addend(2) == 7 * kKvRow, "ELF position 7: row drain at row 7 (word and addend)"); + check(word(108) == (7 * kPtabRow | 0x80000000u) && addend(3) == (7 * kPtabRow | 0x80000000u), + "ELF position 7: record at ptab row 7, folded bit kept (word and addend)"); + check(addend(1) == 0x1E1E0000 && word(56) == 0x1E1E0000, "ELF position 7: the weight fill untouched"); + apply_elf_position(e, sites, 0, kKvRow, kPtabRow); + check(word(8) == kKvRow / 4 && addend(2) == 0 && addend(3) == 0x80000000u, "ELF position 0: one masked row, row 0"); + + // What XRT gets: the assembled full ELF has the fold cleared in the addend and the word. + apply_elf_position(e, sites, 5, kKvRow, kPtabRow); + const Bytes full = assemble_full_elf(e, Bytes(64, 0xAB), "axf5", PdiMode::kNone); + const ElfSection fct = elf_section(full, ".ctrltext.0"), frl = elf_section(full, ".rela.dyn"); + check(rd(full, frl.offset + 12 * 3 + 8) == 5 * kPtabRow && rd(full, fct.offset + 4 * 80) == 5 * kKvRow && rd(full, fct.offset + 4 * 108) == 5 * kPtabRow, + "ELF assembled at position 5: record addend and word without bit 31, drain word 5 rows"); + + // Refused: a relocation naming another argument, and two relocations on one BD. + auto throws = [](auto f) { try { f(); } catch (const std::exception&) { return true; } return false; }; + auto bad = rels; + bad[2].sym = 1; + check(throws([&] { elf_position_sites(make_insts_elf(w, args, bad), kKvRow); }), "ELF: a drain relocation on arg 0 is refused"); + bad = rels; + bad.push_back(rels[2]); + check(throws([&] { elf_position_sites(make_insts_elf(w, args, bad), kKvRow); }), "ELF: two relocations on one BD are refused"); + bad = rels; + bad[3].addend = 0x80000800; + check(throws([&] { elf_position_sites(make_insts_elf(w, args, bad), kKvRow); }), + "ELF: an addend that disagrees with its DDR_PATCH word is refused"); +} + +// A stand-in for the decoder: its "DeltaNet state" is an FNV-1a hash of the tokens fed, +// so a restored snapshot is right exactly when it equals the hash of its prefix. +struct FakeDecoder { + uint64_t h = 1469598103934665603ull; + size_t fed = 0; + void reset() { *this = {}; } + void feed(int t) { + h = (h ^ uint64_t(uint32_t(t))) * 1099511628211ull; + ++fed; + } + static uint64_t of(std::span t) { + FakeDecoder d; + for (int x : t) d.feed(x); + return d.h; + } +}; + +// Session::generate's cache handling (npu/lax.cpp) on the stand-in: returns tokens reused. +struct FakeSession { + FakeDecoder dec; + Snapshots snaps; + std::vector cache, turn; + bool restored_ok = true; + size_t generate(const std::vector& prompt, const std::vector& answer) { + const Plan p = plan(cache, snaps, prompt); + if (p.source == Plan::Snapshot) { + const auto& e = snaps.at(size_t(p.snapshot)); + std::memcpy(&dec.h, e.state.data(), 8); + dec.fed = e.tokens.size(); + cache = e.tokens; + snaps.touch(size_t(p.snapshot)); + restored_ok = restored_ok && dec.h == FakeDecoder::of(cache); + } else if (p.source == Plan::Scratch) { + dec.reset(); + cache.clear(); + } + const size_t reused = cache.size(); + const auto pts = snapshot_points(prompt, cache.size(), turn); + auto s = pts.begin(); + for (size_t i = cache.size(); i < prompt.size(); ++i) { + if (s != pts.end() && *s == i && (++s, !snaps.contains(cache))) + if (auto* e = snaps.put(cache)) { + e->state.resize(8); + std::memcpy(e->state.data(), &dec.h, 8); + } + dec.feed(prompt[i]); + cache.push_back(prompt[i]); + } + // The answer: every generated token but the last is fed (the last only picked). + for (size_t i = 0; i + 1 < answer.size(); ++i) { + dec.feed(answer[i]); + cache.push_back(answer[i]); + } + return reused; + } +}; + +void turns_unit() { + using V = std::vector; + Snapshots s(2); + auto put = [&](const V& t, uint8_t tag) { + auto* e = s.put(t); + e->state.assign(4, tag); + return e; + }; + check(s.best(V{1, 2, 3}) == -1, "snapshots: none kept, no match"); + put({1, 2}, 1); + put({1, 2, 3, 4}, 2); + check(s.best(V{1, 2, 3, 4, 5}) == 1 && s.at(1).state[0] == 2, "snapshots: the longest proper prefix wins"); + check(s.best(V{1, 2, 3, 4}) == 0, "snapshots: an equal prompt takes a shorter one (its last token must run)"); + check(s.best(V{1, 2}) == -1 && s.best(V{1, 3, 3, 4, 5}) == -1, "snapshots: no proper prefix, no match"); + check(s.contains(V{1, 2}) && !s.contains(V{1}), "snapshots: contains is exact"); + s.touch(0); // {1,2} now more recent than {1,2,3,4} + const uint8_t* buf = s.at(1).state.data(); + put({7, 8, 9}, 3); + check(s.size() == 2 && !s.contains(V{1, 2, 3, 4}) && s.contains(V{1, 2}) && s.contains(V{7, 8, 9}), + "snapshots: at capacity the least recently used goes"); + check(s.at(1).state.data() == buf, "snapshots: the evicted entry's state buffer is reused"); + auto* again = s.put(V{1, 2}); + check(s.size() == 2 && again == &s.at(0) && again->state[0] == 1, "snapshots: putting an existing prefix refreshes it"); + check(s.bytes() >= 8, "snapshots: bytes counts the kept states"); + Snapshots none(0); + check(none.put(V{1}) == nullptr && none.best(V{1, 2}) == -1, "snapshots: capacity 0 keeps nothing"); + + Snapshots p(4); + p.put(V{1, 2, 3}); + auto pl = plan(V{1, 2}, p, V{1, 2, 3, 4}); + check(pl.source == Plan::Snapshot && pl.reuse == 3 && pl.snapshot == 0, "plan: a longer snapshot beats the live cache"); + pl = plan(V{1, 2, 3}, p, V{1, 2, 3, 4}); + check(pl.source == Plan::Live && pl.reuse == 3, "plan: on a tie, the live cache (nothing to load)"); + pl = plan(V{1, 2, 3, 4}, p, V{1, 2, 3, 4}); + check(pl.source == Plan::Snapshot && pl.reuse == 3, "plan: the same prompt again restores the snapshot before its last token"); + pl = plan(V{5}, p, V{6, 7}); + check(pl.source == Plan::Scratch && pl.reuse == 0, "plan: nothing shared, from scratch"); + pl = plan(V{}, Snapshots(1), V{6, 7}); + check(pl.source == Plan::Scratch, "plan: an empty cache, from scratch"); + + const V turn = {100}; + check(snapshot_points(V{100, 1, 2, 100, 3, 4}, 0, turn) == std::vector{3, 5}, + "points: before the last turn token and before the last token"); + check(snapshot_points(V{100, 1, 2, 100, 3, 4}, 4, turn) == std::vector{5}, + "points: none before where feeding starts"); + check(snapshot_points(V{100, 1, 2}, 0, turn) == std::vector{2}, "points: a turn token at 0 is no snapshot"); + check(snapshot_points(V{1, 2, 100}, 0, turn) == std::vector{2}, "points: a last-token turn token counts once"); + check(snapshot_points(V{1}, 0, turn).empty(), "points: a one-token prompt has none"); + + // A three-turn chat as a client sends it: each follow-up repeats the messages, but + // renders the earlier answer without the think header the prompt opened it with. + constexpr int S = 100, E = 101, NL = 10, U = 20, A = 21, T = 30, TE = 31; + auto user = [&](V text) { + V v = {S, U, NL}; + v.insert(v.end(), text.begin(), text.end()); + v.insert(v.end(), {E, NL}); + return v; + }; + auto asst = [&](const V& text) { + V v = {S, A, NL}; + v.insert(v.end(), text.begin(), text.end()); + v.insert(v.end(), {E, NL}); + return v; + }; + const V open = {S, A, NL, T, NL, NL, TE, NL, NL}; // the prompt's end, thinking off + const V q1 = {1, 2, 3, 4, 5, 6}, q2 = {7, 8, 9}, q3 = {11, 12, 13, 14}; + const V a1 = {41, 42, 43, 44, E}, a2 = {51, 52, E}, a3 = {61, E}; + auto cat = [](std::initializer_list parts) { + V v; + for (const auto& p : parts) v.insert(v.end(), p.begin(), p.end()); + return v; + }; + auto body = [](const V& a) { return V(a.begin(), a.end() - 1); }; + const V p1 = cat({user(q1), open}); + const V p2 = cat({user(q1), asst(body(a1)), user(q2), open}); + const V p3 = cat({user(q1), asst(body(a1)), user(q2), asst(body(a2)), user(q3), open}); + for (size_t cap : {size_t(0), size_t(4)}) { + FakeSession fs; + fs.snaps = Snapshots(cap); + fs.turn = {S}; + const size_t r1 = fs.generate(p1, a1), r2 = fs.generate(p2, a2), r3 = fs.generate(p3, a3); + const std::string tag = "chat, " + std::to_string(cap) + " snapshots: "; + if (cap == 0) { + check(r1 == 0 && r2 == 0 && r3 == 0, tag + "every follow-up starts over (today's exact-extension only)"); + } else { + check(r1 == 0 && r2 == user(q1).size() && r3 == cat({user(q1), asst(body(a1)), user(q2)}).size(), + tag + "each follow-up reuses everything before the last answer (" + std::to_string(r2) + ", " + + std::to_string(r3) + " tokens)"); + check(fs.restored_ok, tag + "every restored state equals its prefix fed from scratch"); + check(fs.snaps.size() <= cap, tag + "at most the capacity kept"); + } + check(fs.dec.h == FakeDecoder::of(fs.cache), tag + "the device state matches the tokens it records"); + // The same last request again (a retry): restores the snapshot before its last token. + const size_t r4 = fs.generate(p3, a3); + check(cap == 0 ? r4 == 0 : r4 == p3.size() - 1, tag + "a retry reuses " + std::to_string(r4) + " tokens"); + // A raw prompt that extends the live cache still continues it. + const size_t live = fs.cache.size(); + const V ext = cat({fs.cache, V{E, NL}, user(V{9}), open}); + check(fs.generate(ext, a3) == live, tag + "a prompt extending the device's tokens continues them"); + } +} + +void unit() { + check(sha_str("") == "e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855", "sha256 of \"\""); + check(sha_str("abc") == "ba7816bf8f01cfea414140de5dae2223b00361a396177a9cb410ff61f20015ad", "sha256 of \"abc\""); + check(sha_str("abcdbcdecdefdefgefghfghighijhijkijkljklmklmnlmnomnopnopq") == + "248d6a61d20638b8e5c026930c3e6039a33ce45964ff2167f6ecedd419db06c1", + "sha256 of the two-block vector"); + + const auto q8 = q8_chunks(8); + check(sha(q8) == "8647fb97cb1e8d15fbbfeb9e180f948100c7907244fd8d40b6d2d207f447825c", "q8 test chunks (LCG)"); + std::vector q41(8 * kQ41); + for (int c = 0; c < 8; ++c) requant_q8(q8.data() + c * kQ8, q41.data() + c * kQ41); + check(sha(q41) == "a854c6d6d5e6b455e5b8f3398c257b9b283c1b5ee6919ec11d785729dc40ecb7", "q8 -> q4_1 = reference"); + + const auto q4 = q4_chunks(2); + check(sha(q4) == "b70797fdc38cded3eb1220a5d99f5eef76cfb3a90777c0e495c7afdaad2c9c0f", "q4 test chunks (LCG)"); + std::vector q4c(2 * kQ41); + for (int c = 0; c < 2; ++c) q4_0_to_q4_1(q4.data() + c * kQ41, q4c.data() + c * kQ41); + check(sha(q4c) == "46a78bf613286056665d8699b4fd6e4ec637c062bb0b95935562fbcfa1b1b302", "signed q4 -> q4_1 = reference"); + check(is_signed_q4(q4.data()) && !is_signed_q4(q4c.data()), "signed-nibble detection (every min zero)"); + + check(perm_sha(2048, [](size_t c) { return std_src(c, 2048); }) == + "b9125fb145c43a8f1fce01a94677f04d8c529e93aceb61d812c88611a1eda2cd", + "std band, 2048 chunks x 2048 wide"); + check(perm_sha(128, [](size_t c) { return std_src(c, 512); }) == + "6f1d98e68120d1228af04ce8065bd7d3448e025a993cf52b9802a8d940b068ca", + "std band, 128 chunks x 512 wide"); + check(perm_sha(1024, [](size_t c) { return std_src(c, 4096); }) == + "9510e09e730a6d3b904cb53ddfc1680ea268d97aef846eb1b428a64a84e1c355", + "std band, 1024 chunks x 4096 wide"); + check(perm_sha(128, down_src) == "8a4458d562f3f2881bc485dcf5a26d0b6e688869ff197bcffdee7ec71b5dbcb4", + "expert down slice"); + check(perm_sha(32, stripe_src) == "23f038693437780b8fbac60e5bf7e22c328ba308129ea0d43ceda6c22fdab70d", + "expert up/gate stripe"); + check(perm_sha(62080, [](size_t k) { return lmhead_src(k, 2048); }) == + "03075dcdb60acada1095d81c64493328285f9104fed5222475af20c1ff298290", + "q8 lm head supertiles"); + + Config c; + c.rotary_dim = 64; + c.rope_theta = 1e7; + std::vector pt(kPtabBytes); + pack_ptab(c, kMaxContext, pt.data()); + check(sha(pt) == "2a11c3c044d7cd4cd3a07abb7c2b5e40e6fe11b2793c891b44ccf5e28d2ad8f7", "position table, 4096 rows"); + + // A synthetic stream: header, a BD write (register 0x1D000, length word), the KV window + // fill on arg 3 at offset 0 for that register + 4, the row drain at one row with the + // translation bit, the record fill on arg 5, and filler ops of every length. + std::vector w = {0x06030100, 0, 0, 0}; + auto op = [&](std::vector o) { w.insert(w.end(), o.begin(), o.end()); }; + op({0x80, 0, 0, 0}); + op({0x01, 0, 0x1D000, 0, 512, 0, 0, 0, 0, 0, 0, 0}); + op({0x81, 0, 0, 0, 0, 0, 0x1D004, 0, 3, 0, 0, 0}); + op({0x00, 0, 0, 0, 0, 0}); + op({0x81, 0, 0, 0, 0, 0, 0x1D024, 0, 3, 0, uint32_t(kKvRow) | 0x80000000u, 0}); + op({0x03, 0, 0, 0, 0, 0, 0}); + op({0x81, 0, 0, 0, 0, 0, 0x1D044, 0, 5, 0, uint32_t(kPtabRow) | 0x80000000u, 0}); + op({0x81, 0, 0, 0, 0, 0, 0x1D064, 0, 0, 0, 12345, 0}); // a weight fill: untouched + const auto patches = position_patches(w, kKvRow); + check(patches.size() == 4, "four position patches in the synthetic stream"); + apply_position(w, patches, 7, kKvRow, kPtabRow); + check(w[4 + 4 + 4] == 7 * kKvRow / 4, "window length: 7 rows in words"); + check(w[4 + 4 + 12 + 10] == 0, "window offset: row 0"); + check(w[4 + 4 + 12 + 12 + 6 + 10] == (7 * kKvRow | 0x80000000u), "row drain: row 7, translation bit kept"); + check(w[4 + 4 + 12 + 12 + 6 + 12 + 7 + 10] == (7 * kPtabRow | 0x80000000u), "record: ptab row 7"); + check(w[4 + 4 + 12 + 12 + 6 + 12 + 7 + 12 + 10] == 12345, "weight fill untouched"); + apply_position(w, patches, 0, kKvRow, kPtabRow); + check(w[4 + 4 + 4] == kKvRow / 4, "position 0 streams one (masked) row"); + + elf_sites_unit(); + turns_unit(); + + const auto cw = cfg_words(0x1'2345'6780'0000ull); + check(cw[0] == 0x67800000u && cw[1] == 0x12345u && cw[2] == 0x1D21C && cw[3] == 0x1D214 && cw[9] == 0x1D214, + "cfg words: pool address lo/hi, then the eight column queues"); +} + +std::map> read_table(const std::string& path) { + std::ifstream f(path); + if (!f) throw std::runtime_error("cannot read " + path); + std::map> t; + for (std::string line; std::getline(f, line);) { + std::istringstream in(line); + std::string name, digest; + size_t bytes = 0; + if (in >> name >> bytes >> digest) t[name] = {bytes, digest}; + } + return t; +} + +void compare_file(const std::string& path, const std::vector& ours, const std::string& name) { + std::ifstream f(path, std::ios::binary | std::ios::ate); + if (!f) return; + std::vector ref(size_t(f.tellg())); + f.seekg(0); + f.read(reinterpret_cast(ref.data()), std::streamsize(ref.size())); + size_t first = 0; + const size_t n = std::min(ref.size(), ours.size()); + while (first < n && ref[first] == ours[first]) ++first; + check(ref.size() == ours.size() && first == n, + name + " byte for byte vs " + path + (first < n ? " (first difference at byte " + std::to_string(first) + ")" : "")); +} + +int model_mode(const std::string& dir, const std::string& table_path, const std::string& ref, const std::string& only, + const std::string& out, int threads) { + const Model m(dir); + const Config c = Config::from_model(m, dir); + const auto table = read_table(table_path); + std::vector names = {"ptab", "normw", "lmpool"}; + for (int L = 0; L < c.layers; ++L) { + names.push_back("pool_" + std::to_string(L)); + names.push_back("consts_" + std::to_string(L)); + } + if (!only.empty()) { + std::vector keep; + std::stringstream ss(only); + for (std::string s; std::getline(ss, s, ',');) keep.push_back(s); + names = keep; + } + std::atomic next{0}; + auto work = [&] { + for (size_t i; (i = next++) < names.size();) { + const std::string& n = names[i]; + std::vector b; + try { + if (n == "ptab") { + b.resize(kPtabBytes); + pack_ptab(c, kMaxContext, b.data()); + } else if (n == "normw") { + b.resize(kHidden * 2); + pack_norm(m, b.data()); + } else if (n == "lmpool") { + b.resize(kLmPoolBytes); + pack_lmhead(m, b.data()); + } else if (n.rfind("pool_", 0) == 0) { + b.resize(kPoolBytes); + pack_pool(m, c, std::stoi(n.substr(5)), b.data()); + } else if (n.rfind("consts_", 0) == 0) { + const int L = std::stoi(n.substr(7)); + b.resize(consts_bytes(c.kinds[size_t(L)])); + pack_consts(m, c, L, b.data()); + } else throw std::runtime_error("unknown buffer " + n); + } catch (const std::exception& e) { + check(false, n + ": " + e.what()); + continue; + } + const auto it = table.find(n); + const std::string d = sha(b); + if (it == table.end()) check(false, n + ": not in the reference table"); + else + check(it->second.first == b.size() && it->second.second == d, + n + " sha256 " + d + (it->second.second == d ? "" : " != reference " + it->second.second)); + if (!ref.empty()) compare_file(ref + "/" + n + ".bin", b, n); + if (!out.empty()) std::ofstream(out + "/" + n + ".bin", std::ios::binary).write(reinterpret_cast(b.data()), std::streamsize(b.size())); + } + }; + if (threads <= 0) threads = 4; + std::vector ts; + for (int t = 0; t < threads; ++t) ts.emplace_back(work); + for (auto& t : ts) t.join(); + return 0; +} + +int insts_mode(const std::string& insts, const std::string& expect) { + std::ifstream f(insts, std::ios::binary | std::ios::ate); + if (!f) throw std::runtime_error("cannot read " + insts); + std::vector w(size_t(f.tellg()) / 4); + f.seekg(0); + f.read(reinterpret_cast(w.data()), std::streamsize(w.size() * 4)); + const auto p = position_patches(w, kKvRow); + std::ifstream e(expect); + if (!e) throw std::runtime_error("cannot read " + expect); + std::vector want; + for (std::string line; std::getline(e, line);) { + std::istringstream in(line); + size_t word; + int kind; + uint32_t flags; + if (in >> word >> kind >> std::hex >> flags) want.push_back({word, PosPatch::Kind(kind), flags}); + } + bool same = p.size() == want.size(); + for (size_t i = 0; same && i < p.size(); ++i) + same = p[i].word == want[i].word && p[i].kind == want[i].kind && p[i].flags == want[i].flags; + for (const auto& x : p) std::printf(" word %zu kind %d flags %08x\n", x.word, int(x.kind), x.flags); + check(same, "lax_a position patches = the reference harness's table"); + return 0; +} + +void pdi_less_positions(const ElfInputs& in); + +// tsv lines: where (word | reloc), index, kind, flags (hex) +int elf_mode(const std::string& dir, const std::string& expect, const std::string& golden) { + const auto in = ElfInputs::read(dir); + std::ifstream f(expect); + if (!f) throw std::runtime_error("cannot read " + expect); + std::vector want; + for (std::string line; std::getline(f, line);) { + if (line.empty() || line[0] == '#') continue; + std::istringstream l(line); + std::string where; + size_t index; + int kind; + uint32_t flags; + if (l >> where >> index >> kind >> std::hex >> flags) + want.push_back({0, PosPatch::Kind(kind), flags, where == "reloc" ? ElfPosSite::Addend : ElfPosSite::CtrlWord, index}); + } + bool same = in.sites.size() == want.size(); + for (size_t i = 0; same && i < want.size(); ++i) + same = in.sites[i].where == want[i].where && in.sites[i].index == want[i].index && in.sites[i].kind == want[i].kind && + in.sites[i].flags == want[i].flags; + for (const auto& x : in.sites) + std::printf(" %s %zu kind %d flags %08x\n", x.where == ElfPosSite::Addend ? "reloc" : "word ", x.index, int(x.kind), + x.flags); + check(same, "lax_a insts.elf position sites = the investigation's attn_sites"); + pdi_less_positions(in); + if (golden.empty()) return 0; + auto cmp = [&](const Bytes& ours, const std::string& file) { + const std::string p = golden + "/" + file; + if (!std::filesystem::exists(p)) { + check(false, file + " missing in " + golden); + return; + } + const Bytes ref = read_file(p); + size_t first = 0; + while (first < std::min(ours.size(), ref.size()) && ours[first] == ref[first]) ++first; + check(ours == ref, "full ELF " + file + " = the harness's" + + (ours == ref ? "" : " (sizes " + std::to_string(ours.size()) + " / " + std::to_string(ref.size()) + + ", first difference at byte " + std::to_string(first) + ")")); + }; + cmp(init_elf(in), "lax_init.elf"); + cmp(linear_elf(in), "lxf.elf"); + for (size_t p = 0; std::filesystem::exists(golden + "/axf_p" + std::to_string(p) + ".elf"); ++p) + cmp(full_elf(in, p, true), "axf_p" + std::to_string(p) + ".elf"); + cmp(norm_elf(in), "ln.elf"); + cmp(head_elf(in), "lm.elf"); + return 0; +} + +// What the decode adds per position is the same ELF without the PDI (npu/lax_elf.h). +void pdi_less_positions(const ElfInputs& in) { + bool same = true; + for (size_t p : {size_t(0), size_t(1), size_t(2), size_t(4095)}) { + const Bytes a = full_elf(in, p), b = full_elf(in, p, true); + same = same && elf_section(a, ".pdi.1").size == 0; + for (const char* sec : {".ctrltext.0", ".rela.dyn", ".dynsym", ".dynstr", ".symtab", ".strtab", ".note.xrt.UID"}) { + const ElfSection x = elf_section(a, sec), y = elf_section(b, sec); + same = same && x.size == y.size && std::equal(a.begin() + x.offset, a.begin() + x.offset + x.size, b.begin() + y.offset); + } + } + check(same, "position configs: no PDI, control code, relocations and symbols as with it (positions 0, 1, 2, 4095)"); +} + +} // namespace + +int main(int argc, char** argv) { + std::string model, table, ref, only, out, insts, patches, elf, sites, golden; + int threads = 0; + for (int i = 1; i < argc; ++i) { + const std::string a = argv[i]; + auto next = [&]() -> std::string { return i + 1 < argc ? argv[++i] : throw std::runtime_error(a + " needs a value"); }; + if (a == "--model") model = next(); + else if (a == "--sha256") table = next(); + else if (a == "--ref") ref = next(); + else if (a == "--only") only = next(); + else if (a == "--out") out = next(); + else if (a == "--threads") threads = std::stoi(next()); + else if (a == "--insts") insts = next(); + else if (a == "--patches") patches = next(); + else if (a == "--elf") elf = next(); + else if (a == "--sites") sites = next(); + else if (a == "--golden") golden = next(); + else { + std::fprintf(stderr, "unknown argument %s\n", a.c_str()); + return 2; + } + } + try { + if (!model.empty()) model_mode(model, table, ref, only, out, threads); + else if (!insts.empty()) insts_mode(insts, patches); + else if (!elf.empty()) elf_mode(elf, sites, golden); + else unit(); + } catch (const std::exception& e) { + std::printf("FAIL %s\n", e.what()); + return 1; + } + std::printf("%s (%d failure%s)\n", failures ? "FAIL" : "PASS", failures, failures == 1 ? "" : "s"); + return failures ? 1 : 0; +} diff --git a/tests/npu_lax_turns.py b/tests/npu_lax_turns.py new file mode 100755 index 00000000..a5432b04 --- /dev/null +++ b/tests/npu_lax_turns.py @@ -0,0 +1,123 @@ +#!/usr/bin/env python3 +# 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. +"""A three-turn chat against `1bit serve`, as an OpenAI client sends it (docs/npu-lax.md, +"Chat follow-ups"): every request carries the whole history, and each earlier answer goes +back without its think block, as Qwen3.6's chat template renders it. Per turn it prints +the prompt tokens, how many the server reused (usage.prompt_tokens_details.cached_tokens) +and from where, the tokens fed again, and the time to the first streamed token, measured +here from sending the request. + +usage: tests/npu_lax_turns.py [--max-tokens 64] [--thinking] [--no-stream] [--min-reuse] + --no-stream plain requests (for servers without stream_options.include_usage): no time to + first token, only the server's prompt time + --min-reuse exit 1 unless every follow-up reuses the history before its last answer: all + of the previous prompt but its last assistant header (at most 8 tokens) +""" +import argparse +import json +import sys +import time +import urllib.error +import urllib.request + +CONTEXT = ( + "Notes for a trip. The Loire valley lies south-west of Paris. Its chateaux were built from the " + "fifteenth to the eighteenth century, when the French court moved between Amboise, Blois and " + "Chambord. Chambord has 426 rooms and a double-helix staircase often credited to Leonardo da Vinci, " + "who spent his last years at Clos Luce in Amboise. Chenonceau spans the river Cher on a gallery of " + "arches; during the First World War it served as a hospital. Trains from Paris Montparnasse reach " + "Tours in about an hour, and most chateaux are within forty kilometres of it. Bicycles can be hired " + "in Tours, Amboise and Blois, and the Loire a Velo route follows the river for 900 km." +) +TURNS = [ + CONTEXT + "\n\nWhich chateau spans a river, and what was it used for in the First World War?", + "How long is the train from Paris, and from which station?", + "Suggest a two-day cycling plan using these notes. Keep it short.", +] + + +def post_stream(url, body): + req = urllib.request.Request(url + "/v1/chat/completions", data=json.dumps(body).encode(), + headers={"Content-Type": "application/json"}) + t0 = time.monotonic() + ttft, text, usage, timings = None, "", None, {} + with urllib.request.urlopen(req, timeout=600) as r: + for raw in r: + line = raw.decode().strip() + if not line.startswith("data: ") or line == "data: [DONE]": + continue + c = json.loads(line[6:]) + if "error" in c: + raise RuntimeError(c["error"]) + for ch in c.get("choices", []): + d = ch.get("delta", {}).get("content", "") + if d and ttft is None: + ttft = time.monotonic() - t0 + text += d + if c.get("usage"): + usage, timings = c["usage"], c.get("timings", {}) + return text, usage, timings, ttft, time.monotonic() - t0 + + +def post_plain(url, body): + body = dict(body, stream=False) + body.pop("stream_options") + req = urllib.request.Request(url + "/v1/chat/completions", data=json.dumps(body).encode(), + headers={"Content-Type": "application/json"}) + t0 = time.monotonic() + with urllib.request.urlopen(req, timeout=600) as r: + c = json.loads(r.read()) + return c["choices"][0]["message"]["content"], c["usage"], c.get("timings", {}), None, time.monotonic() - t0 + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("url") + ap.add_argument("--max-tokens", type=int, default=64) + ap.add_argument("--thinking", action="store_true") + ap.add_argument("--no-stream", action="store_true") + ap.add_argument("--min-reuse", action="store_true") + a = ap.parse_args() + messages, ok = [], True + for i, q in enumerate(TURNS): + messages.append({"role": "user", "content": q}) + body = {"messages": messages, "max_tokens": a.max_tokens, "stream": True, + "stream_options": {"include_usage": True}, + "chat_template_kwargs": {"enable_thinking": a.thinking}} + try: + text, u, tm, ttft, total = (post_plain if a.no_stream else post_stream)(a.url, body) + except urllib.error.HTTPError as e: + print(f"turn {i + 1}: HTTP {e.code}: {e.read().decode()[:300]}") + return 1 + if u is None: + print(f"turn {i + 1}: no usage chunk (stream_options.include_usage unsupported?)") + return 1 + n, cached = u["prompt_tokens"], u["prompt_tokens_details"]["cached_tokens"] + print(f"turn {i + 1}: {n} prompt tokens, {cached} reused ({tm.get('cache', 'none')}), {n - cached} fed; " + f"first token after {'-' if ttft is None else round(ttft, 2)} s (prompt {tm.get('prompt_ms', 0) / 1000:.2f} s " + f"on the server); {u['completion_tokens']} generated, {total:.1f} s in all") + answer = text.split("")[-1].strip() # the template drops the think block + print(f" {answer[:100]!r}") + if a.min_reuse and i > 0 and cached < prev_n - 8: + ok = False + print(f" FAIL: reused {cached}, less than the previous prompt ({prev_n}) but its assistant header") + prev_n = n + messages.append({"role": "assistant", "content": answer}) + return 0 if ok else 1 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/third_party/OpenFlowLM-Next b/third_party/OpenFlowLM-Next new file mode 160000 index 00000000..2490fa6e --- /dev/null +++ b/third_party/OpenFlowLM-Next @@ -0,0 +1 @@ +Subproject commit 2490fa6e6ddb347bdfdb531a54c1ba74516815b6