npu: rotate activations into a folded PB16i container's Hadamard basis

.gitignore

@@ -23,3 +23,6 @@ __pycache__/
 
 # Mesh runtime state (runs, leases, sessions, artifacts) - never commit
 .pi/mesh/
+
+# ctest run in the source tree (kept out of the PR by the copyright check)
+Testing/

CMakeLists.txt

@@ -430,6 +430,7 @@ if(Python3_Interpreter_FOUND)
     # the HRX hardware gate's judges (docs/hw-gate.md) against a fake server and synthetic logs;
     # the gate itself needs the GPU and runs on Strix Halo (.github/workflows/hw-gate.yml)
     add_test(NAME hw_gate_judges COMMAND ${Python3_EXECUTABLE} ${CMAKE_SOURCE_DIR}/tests/hw_gate_test.py)
+    add_test(NAME q4nx_pb16i COMMAND ${Python3_EXECUTABLE} ${CMAKE_SOURCE_DIR}/tests/q4nx_pb16i_test.py)
 endif()
 
 add_executable(npu_full_elf_test tests/npu_full_elf_test.cpp)
@@ -472,6 +473,13 @@ add_test(NAME npu_blocked_layout
          COMMAND ${CMAKE_SOURCE_DIR}/tests/npu_blocked_layout_test.sh
                  ${CMAKE_CXX_COMPILER} ${CMAKE_SOURCE_DIR}/npu/forward/include
                  ${CMAKE_SOURCE_DIR}/tests/npu_blocked_layout_test.cpp)
+# The pb16i weight image (q4nx_weight_layout "pb16i"): the same kind of contract, plus the
+# packer that builds the images (tools/q4nx_pb16i.py).  No device, no model: the real-GGUF
+# leg is skipped when the model is not on the machine.
+add_test(NAME npu_pb16i_layout
+         COMMAND ${CMAKE_SOURCE_DIR}/tests/npu_pb16i_layout_test.sh
+                 ${CMAKE_CXX_COMPILER} ${CMAKE_SOURCE_DIR}/npu/forward/include
+                 ${CMAKE_SOURCE_DIR}/tests/npu_pb16i_layout_test.cpp)
 # Crafted model.q4nx files the loader and the packer must refuse (written at run time); runs in CI.
 add_executable(npu_model_bounds_test tests/npu_model_bounds_test.cpp)
 target_link_libraries(npu_model_bounds_test PRIVATE onebit_npu)

npu/forward/engine-src/npu_engine_i8ctx_inc.h

@@ -27,24 +27,49 @@
 #include <cmath>
 #include <vector>
 #include <memory>
+#include <climits>
+#include <cstdlib>
+#include <map>
 #include <string>
 
 #include "q4nx_weight_layout.h"   // bpos(): the container's device weight layout
+#include "q4nx_pb16i.h"          // the pb16i image: integer-exact ternary decode weights
 #include <algorithm>
 #include "rt.h"  // the own NPU runtime (ONEBIT_NPU_RT_DIR)
 
 
+// One full ELF, one PDI, one hw_context -- and one kernel handle per instance name
+// ("main:<instance>").  This is not a convenience: switching hw_context costs ~2.7 ms on
+// this box, independent of the stream size (measured 2026-10-31: 0.488 ms per launch back
+// to back in one context, 3.005 alternating two, 3.203 alternating seven), so every design
+// whose ELF carries several instances shares one context and only the kernel handle
+// differs.  The pool below is process-wide because the engine opens one device and one
+// model; it is keyed by the ELF's canonical path, so two roles pointing at the same
+// artifact (G and U share one pb16i ELF today) also share its context.
+struct ElfHost {
+    std::unique_ptr<xrt::elf> elf;
+    std::unique_ptr<xrt::hw_context> hc;
+    std::map<std::string, std::unique_ptr<xrt::ext::kernel>> kernels;
+    xrt::ext::kernel* get(const std::string& inst) {
+        auto it = kernels.find(inst);
+        if (it != kernels.end()) return it->second.get();
+        auto k = std::make_unique<xrt::ext::kernel>(*hc, inst);
+        xrt::ext::kernel* raw = k.get();
+        kernels.emplace(inst, std::move(k));
+        return raw;
+    }
+};
+
 struct I8Ctx {
     int MD, KD, ND, NL;
-    std::unique_ptr<xrt::hw_context> hc;
     // Consume a `aiecc --generate-full-elf` artifact directly: the ELF bundles
     // the PDIs + TXN control code, so there is no instruction file. The
     // hw_context is created from the ELF, and kext is the resulting
     // xrt::ext::kernel whose signature is (A, B, C). Requires the design's
     // runtime sequence to load the ELF's PDI (aiex.npu.load_pdi; see
     // generators/build_zaya_m8.sh), otherwise the runtime never configures the array.
-    std::unique_ptr<xrt::elf> elf;
-    std::unique_ptr<xrt::ext::kernel> kext;
+    std::shared_ptr<ElfHost> host_;
+    xrt::ext::kernel* kext = nullptr;
     std::unique_ptr<xrt::bo> bA, bC;
     std::vector<std::unique_ptr<xrt::bo>> layerB;     // weight BOs
     int8_t* Am;
@@ -61,6 +86,24 @@ struct I8Ctx {
     // Set by Q4nxNpuForward from the container's config before the first pack.
     bool blocked_b = false;
 
+    // Weight image "pb16i" (q4nx_weight_layout "pb16i"): the BO holds a pb16i sidecar
+    // unchanged (codes + int16 multipliers), there is no packing and no int8 scale --
+    // the per-column 2^e_ref table and the int32/int64 mode come from the sidecar's
+    // own `.ref` file, and the host applies the single rounding.  The output BO grows
+    // to two int32 planes for the int64 mode (S = hi*2^21 + lo), which the kernel
+    // writes as [M][2N]; bC is allocated for both modes so one BO serves either.
+    bool pb16i = false;
+    // int32 planes the pb16i ELF leaves its sum in: 1 (plain int32) or 2 (ACC64:
+    // lo = sum of (t & 2^21-1), hi = sum of (t >> 21), S = hi*2^21 + lo).  Both are
+    // exact, so a tensor whose own bound fits int32 may still run on the two-plane
+    // kernel; the reverse is not allowed, and load_pb16i_into refuses it.
+    int pb16i_planes = 1;
+
+    // Bytes of one weight image: the int8 layouts are KD*ND, pb16i is its own size.
+    size_t w_bytes() const {
+        return pb16i ? q4nx_pb16i::w_bytes(KD, ND) : (size_t)KD * (size_t)ND;
+    }
+
     // ── Init: consume a FULL ELF directly ──
     // Kernel signature: (bo0, bo1, bo2) == (A, B, C). No instruction BO: the
     // TXN control stream lives in the ELF's .ctrltext. Kernel name is
@@ -79,16 +122,36 @@ struct I8Ctx {
                             "(K=%d N=%d)\n", KD, ND);
             return false;
         }
+        if (pb16i && !q4nx_pb16i::fits(KD, ND)) {
+            fprintf(stderr, "  I8Ctx::init_elf: the pb16i layout needs K%%256==0 and N%%1024==0 "
+                            "(K=%d N=%d)\n", KD, ND);
+            return false;
+        }
         try {
-            elf = std::make_unique<xrt::elf>(std::string(ep));
-            if (!elf->is_full_elf()) {
-                fprintf(stderr, "  I8Ctx::init_elf: %s is not a full ELF\n", ep);
-                return false;
+            // One context per artifact: a second design on the same ELF reuses it, so
+            // going from one role to another inside that ELF does not pay a switch.
+            char canon[PATH_MAX];
+            const std::string key = realpath(ep, canon) ? std::string(canon) : std::string(ep);
+            static std::map<std::string, std::shared_ptr<ElfHost>> pool;
+            auto it = pool.find(key);
+            if (it != pool.end()) {
+                host_ = it->second;
+            } else {
+                auto h = std::make_shared<ElfHost>();
+                h->elf = std::make_unique<xrt::elf>(key);
+                if (!h->elf->is_full_elf()) {
+                    fprintf(stderr, "  I8Ctx::init_elf: %s is not a full ELF\n", ep);
+                    return false;
+                }
+                fprintf(stderr, "  full ELF ok: uuid=%s\n",
+                        h->elf->get_cfg_uuid().to_string().c_str());
+                h->hc = std::make_unique<xrt::hw_context>(d, *h->elf);
+                host_ = h;
+                pool.emplace(key, h);
             }
-            fprintf(stderr, "  full ELF ok: uuid=%s\n",
-                    elf->get_cfg_uuid().to_string().c_str());
-            hc = std::make_unique<xrt::hw_context>(d, *elf);
-            kext = std::make_unique<xrt::ext::kernel>(*hc, std::string(kname));
+            kext = host_->get(std::string(kname));
+            fprintf(stderr, "  I8Ctx::init_elf %s: %s instance %s (%zu in this context)\n",
+                    ep, key.c_str(), kname, host_->kernels.size());
         } catch (std::exception& ex) {
             fprintf(stderr, "  I8Ctx::init_elf failed: %s\n", ex.what());
             return false;
@@ -101,13 +164,15 @@ struct I8Ctx {
             return std::make_unique<xrt::bo>(b);
         };
         bA = mkbo((size_t)MD * KD);
-        bC = mkbo((size_t)MD * ND * 4);
+        // pb16i int64 mode returns two int32 planes ([M][2N]); allocate for both modes
+        // so the tensor's own mode byte, not the geometry, decides what is used.
+        bC = mkbo((size_t)MD * ND * 4 * (pb16i ? 2 : 1));
         Am = (int8_t*)bA->map();
         Cm = (int32_t*)bC->map();
         layerB.resize(NL);
         group_scales.resize(NL);   // packB() writes group_scales[l]
         for (int l = 0; l < NL; l++)
-            layerB[l] = mkbo((size_t)KD * ND);
+            layerB[l] = mkbo(w_bytes());
         initialized = true;
         return true;
     }
@@ -121,12 +186,80 @@ struct I8Ctx {
     std::unique_ptr<xrt::bo> make_weight_bo(xrt::device& d) {
         // The amdxdna driver rejects a group-less xrt::bo, so use the same
         // xrt::ext::bo allocation init_elf's mkbo uses.
-        xrt::ext::bo b{d, (size_t)KD * ND};
+        xrt::ext::bo b{d, w_bytes()};
         auto bo = std::make_unique<xrt::bo>(b);
-        if (void* m = bo->map()) memset(m, 0, (size_t)KD * ND);
+        if (void* m = bo->map()) memset(m, 0, w_bytes());
         return bo;
     }
 
+    // ── pb16i: the weight image is the sidecar file, copied unchanged ──
+    // `<dir>/<stem>.pb16i` is the image and `<dir>/<stem>.ref` the side table
+    // (e_ref int16[N] + mode byte).  Nothing is requantised, so a mismatch is a
+    // wrong file, not a rounding difference: every check here is loud.
+    bool load_pb16i_into(xrt::bo& bo, const std::string& stem,
+                         std::vector<int16_t>& eref, int& mode) {
+        const size_t want = q4nx_pb16i::w_bytes(KD, ND);
+        const std::string imgp = stem + ".pb16i", refp = stem + ".ref";
+        FILE* f = fopen(imgp.c_str(), "rb");
+        if (!f) { fprintf(stderr, "  pb16i: cannot open %s\n", imgp.c_str()); return false; }
+        std::vector<uint8_t> img(want);
+        const size_t got = fread(img.data(), 1, want, f);
+        const int extra = fgetc(f);
+        fclose(f);
+        if (got != want || extra != EOF) {
+            fprintf(stderr, "  pb16i: %s is not %zu bytes (K=%d N=%d)\n", imgp.c_str(), want,
+                    KD, ND);
+            return false;
+        }
+        f = fopen(refp.c_str(), "rb");
+        if (!f) { fprintf(stderr, "  pb16i: cannot open %s\n", refp.c_str()); return false; }
+        std::vector<uint8_t> ref(q4nx_pb16i::ref_bytes(ND) + 1);
+        const size_t rgot = fread(ref.data(), 1, ref.size(), f);
+        fclose(f);
+        if (rgot != q4nx_pb16i::ref_bytes(ND) ||
+            !q4nx_pb16i::parse_ref(ref.data(), rgot, ND, &eref, &mode)) {
+            fprintf(stderr, "  pb16i: %s is not e_ref[%d] + one mode byte\n", refp.c_str(), ND);
+            return false;
+        }
+        if (bo.size() < want) {
+            fprintf(stderr, "  pb16i: weight BO is %zu B, needs %zu\n", (size_t)bo.size(), want);
+            return false;
+        }
+        memcpy(bo.map(), img.data(), want);
+        bo.sync(XCL_BO_SYNC_BO_TO_DEVICE);
+        // NPU_PB16I_HASH=1: one line per image actually loaded, so a run can answer "is this
+        // tensor's image non-zero, and is it different from its neighbours" without a device
+        // probe.  The stem carries (layer, tag), which is what a per-launch gate cannot see.
+        if (getenv("NPU_PB16I_HASH")) {
+            unsigned long long h = 1469598103934665603ULL;
+            const size_t n = want < 4096 ? want : 4096;
+            for (size_t i = 0; i < n; i++) { h ^= img[i]; h *= 1099511628211ULL; }
+            unsigned long long nz = 0;
+            for (size_t i = 0; i < n; i++) nz += (img[i] != 0);
+            fprintf(stderr, "[pb16i-image] %s mode=%d planes=%d first4k_fnv=%016llx nonzero=%llu/%zu\n",
+                    stem.c_str(), mode, pb16i_planes, h, nz, n);
+        }
+        return true;
+    }
+
+    // C = float32(S * 2^e_ref) * ascale, the only rounding in the chain.  Row 0 is
+    // the decode row; am rows are written, each row being an int32 plane or the
+    // pair [lo | hi] of the int64 mode, exactly as the kernel wrote it, and each
+    // row is dequantised with its own activation scale.
+    void dequant_pb16i(float* out, int am, int an, const float* ascales,
+                       const std::vector<int16_t>& eref, int planes) {
+        const int row = planes > 1 ? 2 * ND : ND;
+        for (int m = 0; m < am; m++) {
+            const int32_t* lo = Cm + (size_t)m * (size_t)row;
+            const float asc = ascales[m];
+            for (int n = 0; n < an; n++) {
+                const int64_t s = q4nx_pb16i::sum_of(lo + n, lo + ND + n, planes > 1 ? 1 : 0);
+                out[(size_t)m * (size_t)an + n] =
+                    q4nx_pb16i::scale_i64(s, eref[(size_t)n]) * asc;
+            }
+        }
+    }
+
     // Cache-friendly per-column int8 quantization into a [KD][ND] padded buffer.
     // The original column-at-a-time loop reads w[i*N + j] with a stride of N
     // floats -- 2 KB at N=512 -- so nearly every element is its own cache line;

npu/forward/include/q4nx_forward.h

@@ -192,10 +192,14 @@ class Q4nxNpuForward {
     bool host_gemm_ = false;
     bool load_layer(int l);
     void npu_gemm(I8Ctx& ctx, const float* A, int K, int N,
-                  const float* Bmat, std::vector<float>& out, int wkey = -1);
+                  const float* Bmat, std::vector<float>& out, int wkey = -1,
+                  const char* sc = nullptr);
     // Batched (M>1) GEMM: M rows in one launch, with per-row activation scales.
     void npu_gemm_rows(I8Ctx& ctx, const float* A, int M, int K, int N,
-                       const float* Bmat, std::vector<float>& out, int wkey = -1);
+                       const float* Bmat, std::vector<float>& out, int wkey = -1,
+                       const char* sc = nullptr);
+    // pb16i sidecar stem for (layer, tensor tag): <pb16i_dir>/blk_<l>_<tag>_weight
+    std::string pb16i_stem(int layer, const char* tag) const;
     // Phase-selected activation tile.  A one-row launch is decode, so it runs
     // this design's M=16 twin (cheap single-row kernel); anything wider is
     // prefill and keeps the M=128 kernel so a chunk still lands in one array
@@ -223,6 +227,7 @@ class Q4nxNpuForward {
     void prof_add(ProfAcc* acc, int design, double pack, double quant, double wait,
                   double rb);
     void prof_dump();
+    bool load_hadamard(const std::string& path);
     // True when this layer's projection `wkey` is already packed in its own BO,
     // so the caller can skip both the Q4NX dequant (load_layer) and the f32->B
     // assembly for the rest of the run.
@@ -266,11 +271,50 @@ class Q4nxNpuForward {
     std::unique_ptr<I8Ctx> ctx16_[Q4NX_MAX_DESIGNS];
     bool m16_ = false;          // at least one role has a working twin
     bool m16_check_ = false;    // NPU_M16_CHECK=1: cross-check every decode launch
+    // Weight image "pb16i" (container config q4nx_weight_layout): the decode family
+    // reads sidecar images instead of packing int8 from the f32 weights, so no
+    // f32 weight buffer is needed for a covered projection at all.  pb16i_dir_ is
+    // where blk_<layer>_<tag>_weight.{pb16i,ref} live (config q4nx_pb16i_dir, env
+    // Q4NX_PB16I_DIR wins).  Empty = the layout is off and nothing changes.
+    bool pb16i_ = false;
+    std::string pb16i_dir_;
+    // A sealed ELF carries every ternary instance in ONE full ELF (one PDI, one
+    // hw_context, 7 kernel instances), so changing phase costs nothing: measured on this
+    // box, a phase change inside the sealed context is 0 ms against ~2.7 ms between two
+    // hw_contexts.  Config `q4nx_pb16i_sealed_elf` names it in the kernels dir; without
+    // it the per-role `_pb16i.elf` files are used as before.
+    std::string pb16i_sealed_;
+    //  [[hadamard]]
+    // The PTQ1 fold: a folded container's weights are W_true rotated by a block-`had_block_`
+    // normalized Hadamard with a per-input-width sign vector, so an activation entering one of
+    // those GEMMs is rotated first, A' = H (A . s).  Config key `q4nx_hadamard` names the binary
+    // sign table (magic "HAD1", u32 count, count x (u32 width, u32 offset), f32 signs).
+    bool hadamard_ = false;
+    int had_block_ = 1024;
+    std::vector<float> had_buf_;
+    double had_ms_ = 0.0;      // host cost of the rotation, accumulated for the profile line
+    long had_n_ = 0;
+    std::unordered_map<int, std::vector<float>> had_signs_;   // input width -> signs
+    // Full-attention QKV as pb16i: the one 14336-wide fused launch has no kernel
+    // (its int64 q needs 448 columns per core), so it runs as the two q halves plus
+    // k and v -- the four launches the arena measured at 1.49 ms per layer.
+    int i_qhalf_ = -1, i_kv_ = -1;
+    // The lm head is not a PTQ1_0 tensor (the container stores it I8), so it has no
+    // pb16i image and keeps an int8 context of its own: the M=16 twin of the widest
+    // K=H design when one is on disk, else the M=128 ELF.  Without this, a pb16i
+    // container ran the head on the M=128 kernel and paid 30 ms a launch (~0.55 s a
+    // token, measured 2026-10-31 08:00).
+    std::unique_ptr<I8Ctx> ctx_head_;
     int32_t m16_worst_ = 0;     // worst int32 maxdiff seen by the check
     long m16_bad_ = 0, m16_checked_ = 0;
     std::vector<int32_t> m16_ref_;   // reference C row, reused per check
     bool prof_ = false;              // NPU_PROFILE=1
     ProfAcc prof_dec_[Q4NX_MAX_DESIGNS], prof_pre_[Q4NX_MAX_DESIGNS];
+    // The lm head has no design of its own (it borrows a K=H geometry and, under pb16i, an
+    // int8 context of its own), so it gets its own row: attributing it to a design put
+    // 15-16 launches of ~27 ms into the G row and made any per-role budget wrong by 0.4 s.
+    ProfAcc prof_head_;
+    int     prof_head_M_ = 0;
     long prof_dec_n_ = 0;            // decode launches (dispatches/token)
     long prof_pre_n_ = 0;
     long prof_tok_n_ = 0;
@@ -363,6 +407,9 @@ class Q4nxNpuForward {
     struct PackedWeight {
         std::unique_ptr<xrt::bo> bo;
         std::vector<float>       scales;
+        // pb16i: the sidecar's host side table.  Empty eref = an int8 pack.
+        std::vector<int16_t>     eref;
+        int                      pb16i_mode = 0;
     };
     std::vector<std::unordered_map<int, PackedWeight>> wcache_;   // [layer][wkey]
     bool weight_cache_ = true;    // NPU_NO_WCACHE=1 disables (A/B)
@@ -414,8 +461,22 @@ class Q4nxNpuForward {
     int lin_khd_ = 0, lin_vhd_ = 0;      // head dims
     int lin_ck_ = 4;                     // depthwise conv kernel
     std::vector<float> wqkvlin_, wz_, woutlin_;          // [CD][H], [VD][H], [H][VD]
-    std::vector<float> wa_, wb_;                          // [nvh][H] alpha/beta
-    std::vector<float> lin_conv_w_, lin_norm_w_, lin_ssm_a_, lin_dt_;  // [CK][CD],[VHD],[nvh],[nvh]
+    // GDN members the HOST reads every token: the alpha/beta projections (the
+    // per-head dot products), the conv window, the gate norm and ssm_a/ssm_dt
+    // (the recurrence).  They are container constants, but the recurrence consumes
+    // a different one per layer, so a single shared buffer forced a re-dequant of
+    // all of them for every layer of every token: 299 ms/token on the 27B, 53% of
+    // all layer host time, for ~2 MB/layer.  Per layer, dequantised once.
+    // NPU_GDN_CACHE=0 restores the per-token reload.
+    struct GdnSmall {
+        std::vector<float> wa, wb;        // [nvh][H]
+        std::vector<float> conv;          // [CK][CD]
+        std::vector<float> norm;          // [VHD]
+        std::vector<float> ssm_a, dt;     // [nvh]
+        bool loaded = false;
+    };
+    std::vector<GdnSmall> gdn_small_;     // [layer]
+    bool gdn_cache_ = true;               // NPU_GDN_CACHE=0 disables
     std::vector<std::vector<float>> conv_state_, rec_state_;  // [CD][CK-1], [nvh][khd][vhd]
     // One GDN token step: xn is the normed input; out gets the out_proj result.
     bool gdn_layer(int l, const std::vector<float>& xn, std::vector<float>& out);