npu: a serve loads the packed lm-head from the image cache instead of rebuilding it (fixes #1716)

CMakeLists.txt

@@ -1343,8 +1343,12 @@ target_link_libraries(npu_derive_designs_dump PRIVATE m)
 # actually calls against the reference; the test links both. It is C++, so CMake switches this target's linker to CXX.
 # q4nx_head_image.c is the image cache's file format (engine#1716's cache slice, this branch): the identity arm above pins
 # what the image must contain, this one pins that what is written and read back is byte-identical to it.
+# q4nx_head_tile_cache.cpp is what a serve calls per head tile (load the image a previous serve wrote, or pack and write
+# it): the arms here are the cache's behaviour -- a second pass does not pack, a changed key does, a damaged or stale image
+# is replaced rather than used -- with a counter on the gather, so "rebuilt on every serve" cannot come back unnoticed.
 add_executable(npu_design_grid_test tests/npu_design_grid_test.c npu/forward/src/q4nx_pack.c npu/forward/src/q4nx_dequant.c
-                                  npu/forward/src/q4nx_pack_blocked.cpp npu/forward/src/q4nx_head_image.c)
+                                  npu/forward/src/q4nx_pack_blocked.cpp npu/forward/src/q4nx_head_image.c
+                                  npu/forward/src/q4nx_head_tile_cache.cpp)
 target_include_directories(npu_design_grid_test PRIVATE npu/forward/include)
 target_link_libraries(npu_design_grid_test PRIVATE m)
 # engine#1779 fault 4 ("Gemma 4 tool-call parsing"): the answer parser is header-only (app/npu_answer_parser.h), so
@@ -1584,6 +1588,9 @@ endif()
         npu/forward/src/q4nx_pack.c
         # The production weight packer, extracted from I8Ctx so it can be tested without the device runtime (#1416).
         npu/forward/src/q4nx_pack_blocked.cpp
+        # The lm-head image cache (engine#1716): the file format, and the load-or-pack a serve calls per head tile.
+        npu/forward/src/q4nx_head_image.c
+        npu/forward/src/q4nx_head_tile_cache.cpp
         npu/forward/src/q4nx_semantics.c
         npu/forward/src/q4nx_forward.cpp
         # The NPU target interface and this box's target (engine#799 P4): the model keeps the maths, the

docs/npu.md

@@ -262,6 +262,34 @@ cache (`$ONEBIT_Q4NX_CACHE`, else `~/.cache/1bit/q4nx`) and hands the loader tha
 device streams its 5,120-byte bf16 q4_1 while the container on disk stays exact. The repack is where
 the bf16 rounding reappears for the device, never in the container.
 
+**The packed lm-head is cached beside it (engine#1716).** A model whose head runs as int8 tiles
+packs those tiles on the host — gather `kpad × T` weights per tile, quantise each column to int8
+with its own scale — and until this cache that happened on every serve. The image is a function
+of the head's stored bytes and of the kernel set (`kpad`, the tile width `T`, the blocked or
+row-major layout), so the forward keeps it in `<Q4NX cache>/head-images/` as
+`headimg-<key digest>-t<tile>.q4hi`, one file per tile. The digest covers the format version,
+`kpad`, `T`, the layout, and a hash of the head tensor's bytes together with the vocabulary size,
+the width and the dtype it is decoded at, so another container, another kernel set or a re-import
+names different files.
+
+A cached tile is used only when both of these hold, otherwise it is packed as before and the file
+is replaced, with one stderr line saying why:
+
+- the file reads back whole: format version, geometry and the hash over its payload
+  (`q4nx_head_image_read`);
+- a sample of its columns (8 per tile, both ends and an even spread) packs to the same scale and
+  the same codes from the weights the forward holds at that moment, through the production packer
+  (`q4nx_pack_B_int8_blocked`). This is the check a payload hash cannot make: an intact image
+  written for other weights is caught here.
+
+The serve log line says what happened: `lm head: N packed weight tile(s) resident (…); H from the
+image cache, P packed (W written for the next serve) in S s`. `ONEBIT_HEAD_IMAGE_CACHE=0` turns
+the cache off (nothing is read or written). The files are `kd × nd` bytes plus `nd` scales per
+tile, so a head costs about `kpad × vocabulary` bytes on disk per kernel set; nothing prunes them,
+the same as the device-form copies next to them. Ternary (`pb16i`) heads and the gemv head do not
+use it: their image already is a file. The behaviour is pinned host-side in
+`tests/npu_design_grid_test.c` with a counter on the gather.
+
 **Q8 has two code layouts under the same chunk size.** Qwen3.5-4B's lm_head (tied to the
 embedding) stores code (row r, column k) at byte `k*32 + r`. Qwen3.6-35B-A3B's Q8 tensors
 (untied lm_head, projections) store it at `(r/16)*4096 + (k/32)*512 + (k%32)*16 + r%16`.

npu/forward/engine-src/npu_engine_i8ctx_inc.h

@@ -514,6 +514,15 @@ struct I8Ctx {
         sout = (float)(ssum / N);
     }
 
+    // Fill an (already-allocated) weight BO through `fill(Bm)`, which writes the KD x ND image itself, then sync it.
+    // packB_into() above is this with the packer as `fill`; the head-image cache (engine#1716) is this with
+    // q4nx_head_tile_image(), which either reads the image a previous serve wrote or packs it the same way.
+    template <class F>
+    void fill_weight_bo(xrt::bo& bo, F&& fill) {
+        fill((int8_t*)bo.map());
+        bo.sync(XCL_BO_SYNC_BO_TO_DEVICE);
+    }
+
     // Async launch with an arbitrary weight BO (resident-expert path).
     inline xrt::run launch_async_with_bo(xrt::bo& wbo, const float* A,
                                          int am, int ak, float ascale) {

npu/forward/include/q4nx_head_image.h

@@ -26,13 +26,14 @@
 //
 // What the hash does NOT prove: that the image was built from THESE weights. An image from another pack with the same kd/nd
 // hashes faithfully to its own bytes and passes every check here, so a stale image left beside a different container of the
-// same head geometry would be accepted. Nothing calls this API yet; whoever first does must key the path on the container's
-// identity (or add a source hash beside the payload hash) before a stale file can be read as this pack's head.
+// same head geometry would be accepted. q4nx_head_tile_image() below is the caller, and it closes that gap twice: the file name carries a digest of the head's
+// own bytes (Q4nxHeadImageKey), and a hit re-packs a sample of columns from the live weights before the image is used.
 //
 // Absence is not an error: a container without an image simply packs as it did before.
 #ifndef Q4NX_HEAD_IMAGE_H
 #define Q4NX_HEAD_IMAGE_H
 
+#include <stddef.h>
 #include <stdint.h>
 
 #ifdef __cplusplus
@@ -87,6 +88,49 @@ uint64_t q4nx_head_image_key_digest(const Q4nxHeadImageKey *key);
 int q4nx_head_image_key_name(const Q4nxHeadImageKey *key, char *out, int out_n);
 
 
+// ── The tile cache itself (engine#1716): what a serve calls instead of packing ────────────────────────────────────────
+//
+// One head tile, filled into `codes` (kd*nd bytes, the weight BO's image) and `scales` (nd floats):
+//
+//   * HIT   `<dir>/headimg-<key digest>-t<tile>.q4hi` reads back whole (version, geometry, payload hash) AND a bounded
+//           sample of its columns equals what the packer produces for those columns from the weights the caller holds
+//           NOW. Then the image is used and the tile is not gathered or packed.
+//   * MISS  anything else: no file, a refused file, or a sample that disagrees. The tile is gathered through `col`,
+//           packed by q4nx_pack_B_int8_blocked() -- the packer the forward has always called -- and written for the next
+//           serve. A refused file is replaced, and `why` says what was wrong with it.
+//
+// The sample is what makes a hit safe rather than plausible. The payload hash proves the file is intact; the key proves
+// it was written for this geometry and these head bytes; the sample proves that the bytes about to be copied into a
+// weight BO are the bytes this build's packer writes for this model's weights, by packing those columns again. It is
+// possible because the packer is column-separable: a column's scale is the max over its own K values and each code is
+// rounded with that column's scale alone, so one column packs to the same bytes alone as inside its tile.
+//
+// `col(ctx, n, out)` writes the `k` values of vocabulary column `n` (rows past the model's width as 0). `n0` is the
+// tile's first vocabulary column and `ncols` how many real columns it holds; the image's remaining columns are zero.
+// `k`/`t` are the packer's logical K/N (the key's kpad and T), `kd`/`nd` the BO's padded geometry.
+//
+// Never fails for a cache reason: with `dir` NULL or unwritable it packs and reports wrote = 0. Returns 1 when `codes` and
+// `scales` hold the tile's image, 0 only for unusable arguments or a packer refusal (nothing usable was produced).
+typedef void (*Q4nxHeadColFn)(void *ctx, int n, float *out_k);
+
+typedef struct {
+    int hit;       // 1: the image came from the cache and its sample verified
+    int wrote;     // 1: the image was packed and written for the next serve
+    int refused;   // 1: a cached file was present and NOT used; `why` names the reason
+    int sampled;   // columns re-packed to verify a hit
+    char why[320];
+} Q4nxHeadTileResult;
+
+// Columns re-packed per tile to verify a cached image before it is used.
+#define Q4NX_HEAD_TILE_SAMPLE 8
+
+int q4nx_head_tile_image(const char *dir, const Q4nxHeadImageKey *key, int tile, int n0, int ncols, int k, int t, int kd,
+                         int nd, Q4nxHeadColFn col, void *ctx, int8_t *codes, float *scales, Q4nxHeadTileResult *res);
+
+// The key's `head_hash`: FNV-1a over the head's stored bytes, then over `extra` (the facts that change what those bytes
+// decode to without changing them -- the caller lists them). 0 for no data.
+uint64_t q4nx_head_source_hash(const void *data, size_t nbytes, const uint32_t *extra, int n_extra);
+
 #ifdef __cplusplus
 }
 #endif

npu/forward/src/q4nx_forward.cpp

@@ -37,6 +37,7 @@
 #include "residual-dump.h"
 #include <cstdlib>
 #include <cstring>
+#include <filesystem>
 #include <fstream>
 
 #include "rt.h"  // the own NPU runtime (ONEBIT_NPU_RT_DIR)
@@ -48,6 +49,7 @@
 #include "q4nx_device.h"  // an imported q4nx-2 container's device form (engine#1185)
 #include "q4nx_semantics.h"
 #include "q4nx_pack.h"
+#include "q4nx_head_image.h"   // the lm-head tile cache (engine#1716)
 
 #if defined(_OPENMP)
 #include <omp.h>
@@ -5406,7 +5408,6 @@ bool Q4nxNpuForward::run_lm_head(const std::vector<float>& h, std::vector<float>
     if (head_packed &&
         ((int)lm_head_packed_.size() != ntiles || lm_head_packed_T_ != T)) {
         lm_head_packed_.clear();
-        std::vector<float> bt(tile_elems);
         if (head_pb16i) {
             // The image is the file: no dequant, no pack, nothing to quantise on the host.  A
             // tile whose side table asks for the two-plane accumulator is refused, not guessed:
@@ -5433,20 +5434,76 @@ bool Q4nxNpuForward::run_lm_head(const std::vector<float>& h, std::vector<float>
             fprintf(stderr, "  [forward] lm head: %d ternary image(s) resident (H=%d T=%d), "
                             "f32 head released\n", ntiles, H, T);
         } else {
+        // engine#1716: this loop used to gather and pack every tile on EVERY serve (10-17 s on a quiet box for the
+        // h4096 shapes, minutes on a contended one). The image is a pure function of the head's stored bytes and the
+        // kernel set's geometry, so a serve now asks q4nx_head_tile_image() for it: a previous serve's file is used
+        // when it reads back whole AND a sample of its columns re-packs to the same bytes from the weights held here;
+        // anything else packs exactly as before (same packer, same padded KD x ND image, same per-column scales as
+        // the per-token packB below) and leaves the file for the next serve. ONEBIT_HEAD_IMAGE_CACHE=0 turns it off.
+        const auto head_t0 = std::chrono::steady_clock::now();
+        std::string head_dir;
+        Q4nxHeadImageKey head_key = {(uint32_t)kpad_, (uint32_t)T, lctx.blocked_b ? 1u : 0u, 0};
+        {
+            const char* off = getenv("ONEBIT_HEAD_IMAGE_CACHE");
+            const bool cache_on = !(off && off[0] == '0' && off[1] == 0);
+            // The head's stored bytes: the separate lm_head.weight of an untied model, else the token table lm_at()
+            // falls back to (which a tied container may itself store as lm_head.weight).
+            TensorDesc* src = !lm_head_f_.empty() ? model_tensor_by_name(mw_, "lm_head.weight") : nullptr;
+            if (!src) {
+                src = &mw_->embed_tokens;
+                if (src->data_size == 0 || src->name[0] == '\0') src = model_tensor_by_name(mw_, "lm_head.weight");
+            }
+            const void* sp = (cache_on && src) ? model_tensor_data(mw_, src) : nullptr;
+            // An image the layout cannot hold as kd x nd bytes of T columns is not cached: it packs as before.
+            if (sp && src->data_size > 0 && (int)lctx.ND == T && (int)lctx.KD >= kpad_) {
+                // What changes the decoded head without changing its bytes: the vocabulary, the width it is read at,
+                // whether it is the tied table, and the dtype the decoder dispatches on.
+                uint32_t extra[8] = {(uint32_t)NV, (uint32_t)H, (uint32_t)lm_w_, lm_head_f_.empty() ? 1u : 0u, 0, 0, 0, 0};
+                std::memcpy(&extra[4], src->dtype, std::min(sizeof src->dtype, (size_t)16));
+                head_key.head_hash = q4nx_head_source_hash(sp, (size_t)src->data_size, extra, 8);
+                try {
+                    head_dir = onebit::npu::q4nx::cache_root() + "/head-images";
+                    std::error_code ec;
+                    std::filesystem::create_directories(head_dir, ec);
+                    if (ec) head_dir.clear();
+                } catch (const std::exception&) {
+                    head_dir.clear();   // no home and no ONEBIT_Q4NX_CACHE: pack, as before
+                }
+            }
+        }
+        struct HeadCol {
+            const Q4nxNpuForward* self;
+            int k;
+        } head_col = {this, kpad_};
+        const Q4nxHeadColFn head_col_fn = [](void* ctx, int n, float* out) {
+            const HeadCol* c = (const HeadCol*)ctx;
+            for (int k = 0; k < c->k; k++) out[k] = c->self->lm_at(n, k);
+        };
+        int head_hits = 0, head_wrote = 0;
         for (int t = 0; t < ntiles; t++) {
             const int n0 = t * T, n1 = std::min(NV, n0 + T);
-            std::fill(bt.begin(), bt.end(), 0.0f);   // padded columns must be zero
-            for (int n = n0; n < n1; n++)
-                for (int k = 0; k < kpad_; k++) bt[(size_t)k * T + (n - n0)] = lm_at(n, k);
             PackedWeight pw;
             pw.bo = lctx.make_weight_bo(*dev_);
-            float sout = 1.0f;
-            // Same packer, same padded KD x ND BO, same per-column scales as the
-            // per-token packB below, so the packed bytes are bit-identical.
-            lctx.packB_into(*pw.bo, bt.data(), kpad_, T, sout, pw.scales);
+            pw.scales.assign((size_t)lctx.ND > (size_t)T ? (size_t)lctx.ND : (size_t)T, 0.0f);
+            Q4nxHeadTileResult hres;
+            bool ok = false;
+            lctx.fill_weight_bo(*pw.bo, [&](int8_t* Bm) {
+                ok = q4nx_head_tile_image(head_dir.empty() ? nullptr : head_dir.c_str(), &head_key, t, n0, n1 - n0, kpad_,
+                                          T, lctx.KD, lctx.ND, head_col_fn, &head_col, Bm, pw.scales.data(), &hres) == 1;
+            });
+            if (!ok) {
+                err_ = std::string("lm head: ") + hres.why;
+                return false;
+            }
+            if (hres.refused)
+                fprintf(stderr, "  [forward] lm head: tile %d's cached image was NOT used and is replaced: %s\n", t, hres.why);
+            head_hits += hres.hit;
+            head_wrote += hres.wrote;
+            pw.scales.resize((size_t)T);
             lm_head_packed_.push_back(std::move(pw));
         }
         lm_head_packed_T_ = T;
+        const double head_s = std::chrono::duration<double>(std::chrono::steady_clock::now() - head_t0).count();
         // Both f32 copies are now dead on this path: lm_head_tiles_ was only ever
         // read by the memcpy below, and lm_at() is reachable only from the
         // host_gemm_ branch (which returns above) and from the gather just done.
@@ -5455,8 +5512,9 @@ bool Q4nxNpuForward::run_lm_head(const std::vector<float>& h, std::vector<float>
         lm_head_f_.clear();
         lm_head_f_.shrink_to_fit();
         fprintf(stderr, "  [forward] lm head: %d packed weight tile(s) resident "
-                        "(H=%d T=%d), f32 head + f32 tile gather released\n",
-                ntiles, H, T);
+                        "(H=%d T=%d), f32 head + f32 tile gather released; %d from the image cache, %d packed "
+                        "(%d written for the next serve) in %.2f s\n",
+                ntiles, H, T, head_hits, ntiles - head_hits, head_wrote, head_s);
         }
     }
     std::vector<float> Bt;