diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 8546c451..cc062e68 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 tokenizer_test -j + run: cmake --build build --target onebit npu_full_elf_test npu_pack_test npu_q4nx_test tokenizer_test -j - name: Test run: ctest --test-dir build --output-on-failure diff --git a/CMakeLists.txt b/CMakeLists.txt index 36494422..b9d7f8a7 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -161,7 +161,7 @@ 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/tokenizer.cpp) +add_library(onebit_npu STATIC npu/full_elf.cpp npu/model.cpp npu/pack.cpp npu/q4nx.cpp npu/tokenizer.cpp) target_include_directories(onebit_npu PUBLIC npu) target_link_libraries(onebit_npu PUBLIC nlohmann_json::nlohmann_json PRIVATE PkgConfig::PCRE2) @@ -177,6 +177,10 @@ if(ONEBIT_NPU_CAPTURED AND ONEBIT_NPU_PDI AND ONEBIT_NPU_GOLDEN) COMMAND npu_full_elf_test ${ONEBIT_NPU_CAPTURED} ${ONEBIT_NPU_PDI} ${ONEBIT_NPU_GOLDEN}) endif() +# Q4NX chunk decoders against 1bit-MONSTER's verified decoders, on committed real tiles +add_executable(npu_q4nx_test tests/npu_q4nx_test.cpp) +target_link_libraries(npu_q4nx_test PRIVATE onebit_npu) +add_test(NAME npu_q4nx COMMAND npu_q4nx_test ${CMAKE_SOURCE_DIR}/tests/golden/q4nx) add_executable(npu_pack_test tests/npu_pack_test.cpp) target_link_libraries(npu_pack_test PRIVATE onebit_npu) set(ONEBIT_NPU_MODEL "" CACHE PATH "Q4NX model directory (model.q4nx + config.json)") diff --git a/docs/npu.md b/docs/npu.md index 19bd735a..28edd439 100644 --- a/docs/npu.md +++ b/docs/npu.md @@ -34,7 +34,7 @@ Step 3 lands in three parts: ``` / - model.q4nx weights (safetensors layout, I8 tiles of 5120 B) + model.q4nx weights (safetensors layout, Q4NX tiles: "Q4NX chunks" below) config.json Hugging Face config: dimensions, rope_theta, eos_token_id tokenizer.json Hugging Face tokenizer npu/layer_ctx1.elf layer kernel, instruction ELFs for context lengths 1, 2, 17 @@ -44,6 +44,24 @@ Step 3 lands in three parts: npu/layer.pdi the design both kernels run on ``` +### Q4NX chunks + +A Q4NX weight is a grid of 32-row x 256-column tiles, one chunk each, row-major over +the grid. The chunk kind is the last dimension of the tensor's shape: `[tiles, 5120]`, +or `[tile rows, tile cols, chunk bytes]` (`npu/q4nx.h` has the byte layouts). + +| Chunk | Kind | Weight | Seen in | +|---|---|---|---| +| 5120 B | q4_1: bf16 scale and zero per 32 columns of a row | `code * scale + zero` | Qwen3, Llama, Gemma, Phi models | +| 4736 B | Q4_K: u8 scale and min per 32 columns, bf16 `S`, `M` per row | `S * scale * code + M * min` | Qwen3.5-4B projections | +| 8704 B | Q8: bf16 scale per 32 columns, int8 codes | `d * code` | Qwen3.5-4B `lm_head` and embedding | + +`npu/q4nx.h` decodes all three, and repacks Q4_K into q4_1 (exact apart from rounding +`S * scale` and `M * min` to bf16), so the lane and dx, which read q4_1, run Q4_K +weights unchanged. Q8 does not fit q4_1. `tests/npu_q4nx_test.cpp` checks the decoders +bit for bit against 1bit-MONSTER's verified decoders on real Qwen3.5-4B and Qwen3-0.6B +tiles (committed in `tests/golden/q4nx`), and the repack against its rounding bound. + `1bit serve -m ` serves such a directory on the NPU behind the OpenAI-compatible API ([serve.md](serve.md)); inside Lemonade, that is what its onebit backend runs (`--onebit npu`). diff --git a/npu/model.cpp b/npu/model.cpp index e0950499..20150427 100644 --- a/npu/model.cpp +++ b/npu/model.cpp @@ -28,6 +28,8 @@ namespace onebit::npu { int Tensor::tiles() const { + // [tile rows, tile cols, chunk bytes] (Q4NX with Q4_K or Q8 chunks, npu/q4nx.h) + if (shape.size() == 3 && shape[0] > 0 && shape[1] > 0) return int(shape[0] * shape[1]); if (shape.size() != 2 || shape[0] <= 0) return 0; // A row is one tile only when it is 5120 bytes wide; narrower rows pack // several to a tile. diff --git a/npu/q4nx.cpp b/npu/q4nx.cpp new file mode 100644 index 00000000..af636eec --- /dev/null +++ b/npu/q4nx.cpp @@ -0,0 +1,115 @@ +// 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 "q4nx.h" + +#include +#include +#include + +namespace onebit::npu::q4nx { + +namespace { + +float bf16(const uint8_t* p) { + const uint32_t u = uint32_t(p[0] | (p[1] << 8)) << 16; + float f; + std::memcpy(&f, &u, 4); + return f; +} + +uint16_t to_bf16(float f) { // round to nearest even + uint32_t u; + std::memcpy(&u, &f, 4); + return uint16_t((u + 0x7FFF + ((u >> 16) & 1)) >> 16); +} + +void put16(uint8_t* p, uint16_t v) { + p[0] = uint8_t(v); + p[1] = uint8_t(v >> 8); +} + +// the q4_1 code byte and nibble of row r, column c +size_t q41_code_at(int r, int c) { return 1024 + size_t((r / 16) * 8 + c / kGroup) * 256 + 8 * (c % kGroup) + (r % 16) / 2; } + +} // namespace + +Chunk chunk_of(const Tensor& t) { + const int64_t last = t.shape.empty() ? 0 : t.shape.back(); + if ((t.shape.size() == 2 || t.shape.size() == 3) && + (last == int64_t(Chunk::Q4_1) || last == int64_t(Chunk::Q4K) || last == int64_t(Chunk::Q8))) + return Chunk(last); + throw std::runtime_error("not a Q4NX chunked tensor (last dimension " + std::to_string(last) + ")"); +} + +void dequant_tile(Chunk kind, const uint8_t* t, float* out, int ld) { + for (int r = 0; r < kRows; r++) { + for (int c = 0; c < kCols; c++) { + const int g = c / kGroup, i = g * kRows + r; + float w; + if (kind == Chunk::Q4_1) { + const uint8_t b = t[q41_code_at(r, c)]; + const int code = (r % 2 == 0) ? (b & 15) : (b >> 4); + w = float(code) * bf16(t + 2 * i) + bf16(t + 512 + 2 * i); + } else if (kind == Chunk::Q4K) { + const uint8_t b = t[512 + c * 16 + r / 2]; + const int code = (r % 2 == 0) ? (b & 15) : (b >> 4); + w = bf16(t + 4608 + 2 * r) * float(t[i]) * float(code) + bf16(t + 4672 + 2 * r) * float(t[256 + i]); + } else { + w = bf16(t + 2 * i) * float(int8_t(t[512 + c * kRows + r])); + } + out[size_t(r) * ld + c] = w; + } + } +} + +void dequant(const Tensor& t, int rows, int cols, float* out) { + const Chunk kind = chunk_of(t); + if (rows % kRows || cols % kCols) throw std::runtime_error("Q4NX tensor size is not whole tiles"); + const int tr = rows / kRows, tc = cols / kCols; + if (t.bytes < uint64_t(tr) * tc * uint64_t(kind)) throw std::runtime_error("Q4NX tensor is shorter than its tiles"); + for (int i = 0; i < tr; i++) + for (int j = 0; j < tc; j++) + dequant_tile(kind, t.data + (size_t(i) * tc + j) * size_t(kind), out + size_t(i) * kRows * cols + size_t(j) * kCols, + cols); +} + +void q4k_to_q4_1(const uint8_t* k, uint8_t* q) { + std::memset(q, 0, size_t(Chunk::Q4_1)); + for (int r = 0; r < kRows; r++) { + const float S = bf16(k + 4608 + 2 * r), M = bf16(k + 4672 + 2 * r); + for (int g = 0; g < kCols / kGroup; g++) { + const int i = g * kRows + r; + put16(q + 2 * i, to_bf16(S * float(k[i]))); + put16(q + 512 + 2 * i, to_bf16(M * float(k[256 + i]))); + } + for (int c = 0; c < kCols; c++) { + const uint8_t b = k[512 + c * 16 + r / 2]; + const uint8_t code = (r % 2 == 0) ? (b & 15) : (b >> 4); + q[q41_code_at(r, c)] |= (r % 2 == 0) ? code : uint8_t(code << 4); + } + } +} + +std::vector as_q4_1(const Tensor& t) { + const Chunk kind = chunk_of(t); + const size_t n = size_t(t.bytes / uint64_t(kind)); + if (kind == Chunk::Q4_1) return std::vector(t.data, t.data + n * size_t(kind)); + if (kind == Chunk::Q8) throw std::runtime_error("a Q8 tensor does not fit q4_1 (8-bit codes)"); + std::vector out(n * size_t(Chunk::Q4_1)); + for (size_t i = 0; i < n; i++) q4k_to_q4_1(t.data + i * size_t(kind), out.data() + i * size_t(Chunk::Q4_1)); + return out; +} + +} // namespace onebit::npu::q4nx diff --git a/npu/q4nx.h b/npu/q4nx.h new file mode 100644 index 00000000..a53cac64 --- /dev/null +++ b/npu/q4nx.h @@ -0,0 +1,61 @@ +// 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 three Q4NX chunk kinds, decoded. A Q4NX weight is a grid of 32-row x 256-column +// tiles, each one chunk; a tensor's chunk kind is the last dimension of its shape +// ([tiles, 5120] or [tile rows, tile cols, chunk bytes]), and tiles are stored row-major +// over the grid. Index g*32 + r means column group g (32 columns) of row r. +// +// q4_1 5120 B bf16 scale[g*32+r] @0, bf16 zero[g*32+r] @512, then 4-bit codes in 16 +// regions of 256 B: region h*8+g (h = r/16), byte 8*j + (r%16)/2, even +// row in the low nibble. w = code * scale + zero +// Q4_K 4736 B u8 scale[g*32+r] @0, u8 min[g*32+r] @256, 4-bit codes @512 at byte +// k*16 + r/2 (even row low), bf16 S[r] @4608, bf16 M[r] @4672. +// w = S[r]*scale*code + M[r]*min +// Q8 8704 B bf16 d[g*32+r] @0, int8 codes @512 at byte k*32 + r. w = d * code +// +// Checked bit for bit against 1bit-MONSTER's verified decoders (npu-infer/tools/ +// q4nx_dequant.py) on real Qwen3.5-4B and Qwen3-0.6B tiles: tests/npu_q4nx_test.cpp. +#pragma once + +#include "model.h" + +#include +#include + +namespace onebit::npu::q4nx { + +constexpr int kRows = 32, kCols = 256, kGroup = 32; + +enum class Chunk : uint32_t { Q4_1 = 5120, Q4K = 4736, Q8 = 8704 }; + +// The chunk kind of a quantized tensor; throws if its shape names none of the three. +Chunk chunk_of(const Tensor& t); + +// One tile to f32: out[r * ld + c] for r < 32, c < 256. +void dequant_tile(Chunk kind, const uint8_t* tile, float* out, int ld); + +// A whole rows x cols tensor to f32 (rows % 32 == 0, cols % 256 == 0). +void dequant(const Tensor& t, int rows, int cols, float* out); + +// One Q4_K tile as a q4_1 tile, so the lane and dx run it unchanged. Exact apart from +// rounding the products S[r]*scale and M[r]*min to bf16 (nearest even). +void q4k_to_q4_1(const uint8_t* q4k, uint8_t* q41); + +// The tensor's tiles as q4_1: a copy for q4_1, converted for Q4_K. Throws for Q8, whose +// 8-bit codes do not fit q4_1. +std::vector as_q4_1(const Tensor& t); + +} // namespace onebit::npu::q4nx diff --git a/tests/golden/q4nx/q4_1.bin b/tests/golden/q4nx/q4_1.bin new file mode 100644 index 00000000..5ec58318 Binary files /dev/null and b/tests/golden/q4nx/q4_1.bin differ diff --git a/tests/golden/q4nx/q4k.bin b/tests/golden/q4nx/q4k.bin new file mode 100644 index 00000000..99756b09 Binary files /dev/null and b/tests/golden/q4nx/q4k.bin differ diff --git a/tests/golden/q4nx/q8.bin b/tests/golden/q4nx/q8.bin new file mode 100644 index 00000000..b6f2178f Binary files /dev/null and b/tests/golden/q4nx/q8.bin differ diff --git a/tests/npu_q4nx_test.cpp b/tests/npu_q4nx_test.cpp new file mode 100644 index 00000000..924c9437 --- /dev/null +++ b/tests/npu_q4nx_test.cpp @@ -0,0 +1,138 @@ +// 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_q4nx_test +// +// The Q4NX chunk decoders (npu/q4nx.h) against 1bit-MONSTER's verified decoders +// (npu-infer/tools/q4nx_dequant.py at 35fafde8f, itself checked against the checkpoint's +// bf16 twin): one real tile of each kind, from Qwen3.5-4B (layer 0 qkv_proj tile 0, Q4_K; +// lm_head tile 0, Q8) and Qwen3-0.6B (layer 0 q_proj tile 0, q4_1). The expected output is +// the SHA-256 of the reference's 32 x 256 f32, so any bit that differs fails. +// Then Q4_K -> q4_1: the converted tile must decode within the bf16 rounding of its two +// scale products (half an ulp of S*scale times the largest code, plus half an ulp of M*min). +#include "../npu/q4nx.h" + +#include +#include +#include +#include +#include +#include +#include + +namespace { + +using onebit::npu::q4nx::Chunk; + +std::vector slurp(const std::string& p) { + std::ifstream f(p, std::ios::binary); + return {std::istreambuf_iterator(f), {}}; +} + +// SHA-256 (FIPS 180-4), enough for one 32 KiB buffer +std::string sha256(const uint8_t* data, size_t len) { + static const 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}; + std::vector m(data, data + len); + m.push_back(0x80); + while (m.size() % 64 != 56) m.push_back(0); + for (int i = 7; i >= 0; i--) m.push_back(uint8_t((uint64_t(len) * 8) >> (8 * i))); + auto rotr = [](uint32_t x, int n) { return (x >> n) | (x << (32 - n)); }; + for (size_t off = 0; off < m.size(); off += 64) { + uint32_t w[64]; + for (int i = 0; i < 16; i++) + w[i] = uint32_t(m[off + 4 * i]) << 24 | uint32_t(m[off + 4 * i + 1]) << 16 | uint32_t(m[off + 4 * i + 2]) << 8 | m[off + 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; + } + char hex[65]; + for (int i = 0; i < 8; i++) std::snprintf(hex + 8 * i, 9, "%08x", h[i]); + return hex; +} + +float bf16(const uint8_t* p) { + const uint32_t u = uint32_t(p[0] | (p[1] << 8)) << 16; + float f; + std::memcpy(&f, &u, 4); + return f; +} + +} // namespace + +int main(int argc, char** argv) { + if (argc < 2) { std::fprintf(stderr, "usage: %s \n", argv[0]); return 2; } + const std::string dir = argv[1]; + struct Case { const char* file; Chunk kind; const char* want; }; + const Case cases[] = { + {"q4k.bin", Chunk::Q4K, "d2fc854e3b61bc0c77004724fefd441af57408519e8113bde87b63fd2ee5c4f5"}, + {"q8.bin", Chunk::Q8, "09ad9e7a61cad612ffcb9cc2a1d7dfd27b6813a6d184329404045d4691c5fdf2"}, + {"q4_1.bin", Chunk::Q4_1, "9f58d4201c4844abc937204c2b7f8d364ff44bc574b0bb87292928c54447d2c8"}, + }; + bool ok = true; + std::vector w(32 * 256), v(32 * 256); + std::vector q4k; + for (const Case& c : cases) { + const auto tile = slurp(dir + "/" + c.file); + if (tile.size() != size_t(c.kind)) { std::printf("FAIL: %s is %zu bytes\n", c.file, tile.size()); return 1; } + onebit::npu::q4nx::dequant_tile(c.kind, tile.data(), w.data(), 256); + const std::string got = sha256(reinterpret_cast(w.data()), w.size() * 4); + const bool same = got == c.want; + ok &= same; + std::printf("%s %-8s w[0][0] %.7g, w[31][255] %.7g\n", same ? "ok " : "FAIL", c.file, w[0], w[32 * 256 - 1]); + if (!same) std::printf(" sha256 %s, want %s\n", got.c_str(), c.want); + if (c.kind == Chunk::Q4K) q4k = tile; + } + + // Q4_K -> q4_1 + std::vector q41(size_t(Chunk::Q4_1)); + onebit::npu::q4nx::q4k_to_q4_1(q4k.data(), q41.data()); + onebit::npu::q4nx::dequant_tile(Chunk::Q4K, q4k.data(), w.data(), 256); + onebit::npu::q4nx::dequant_tile(Chunk::Q4_1, q41.data(), v.data(), 256); + double worst = 0, worst_err = 0; + for (int r = 0; r < 32; r++) { + const float S = bf16(q4k.data() + 4608 + 2 * r), M = bf16(q4k.data() + 4672 + 2 * r); + for (int c = 0; c < 256; c++) { + const int i = (c / 32) * 32 + r; + // half an ulp of each bf16 product: 2^(exponent - 8) + auto half_ulp = [](double x) { return x == 0 ? 0.0 : std::ldexp(1.0, std::ilogb(std::fabs(x)) - 8); }; + const double bound = 15 * half_ulp(double(S) * q4k[i]) + half_ulp(double(M) * q4k[256 + i]) + 1e-12; + const double err = std::fabs(double(v[r * 256 + c]) - w[r * 256 + c]); + worst = std::fmax(worst, err / bound); + worst_err = std::fmax(worst_err, err); + } + } + const bool conv = worst <= 1.0 + 1e-6; + ok &= conv; + std::printf("%s Q4_K -> q4_1: max |error| %.3g, at most %.2f of the bf16 rounding bound\n", conv ? "ok " : "FAIL", worst_err, worst); + std::printf("%s\n", ok ? "PASS" : "FAIL"); + return ok ? 0 : 1; +} diff --git a/tools/copyright.py b/tools/copyright.py index 20ef39ea..379df349 100644 --- a/tools/copyright.py +++ b/tools/copyright.py @@ -36,7 +36,7 @@ ] BLOCK_STYLES = [(("*.md", "*.html", "*.svg"), ("")), (("*.css", "*.js"), ("/*", "*/"))] -EXEMPT = ["third_party/*", "patches/*", "LICENSE", "NOTICE", "*.gguf", "*.json", "*.tsv", "tests/golden/*/tokens.txt", "*/Cargo.lock", "*.png", "*.ico"] +EXEMPT = ["third_party/*", "patches/*", "LICENSE", "NOTICE", "*.gguf", "tests/golden/*.bin", "tests/golden/*/*.bin", "*.json", "*.tsv", "tests/golden/*/tokens.txt", "*/Cargo.lock", "*.png", "*.ico"] def _match(name, patterns):