NPU: read all three Q4NX chunk kinds (q4_1, Q4_K, Q8), repack Q4_K into q4_1

.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

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)")

docs/npu.md

@@ -34,7 +34,7 @@ Step 3 lands in three parts:
 
 ```
 <model dir>/
-  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 <model dir>` 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`).

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.

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 <cstring>
+#include <stdexcept>
+#include <string>
+
+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<uint8_t> 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<uint8_t>(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<uint8_t> 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

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 <cstdint>
+#include <vector>
+
+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<uint8_t> as_q4_1(const Tensor& t);
+
+}  // namespace onebit::npu::q4nx

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 <dir with q4_1.bin q4k.bin q8.bin>
+//
+// 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 <cmath>
+#include <cstdint>
+#include <cstdio>
+#include <cstring>
+#include <fstream>
+#include <string>
+#include <vector>
+
+namespace {
+
+using onebit::npu::q4nx::Chunk;
+
+std::vector<uint8_t> slurp(const std::string& p) {
+    std::ifstream f(p, std::ios::binary);
+    return {std::istreambuf_iterator<char>(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<uint8_t> 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 <golden dir>\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<float> w(32 * 256), v(32 * 256);
+    std::vector<uint8_t> 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<const uint8_t*>(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<uint8_t> 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;
+}

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):