From d471d07a4e9213c45aa1dafd4c9581d377aa85f1 Mon Sep 17 00:00:00 2001 From: bong-water-water-bong Date: Tue, 10 Nov 2026 10:46:22 -0400 Subject: [PATCH] npu: a serve loads the packed lm-head from the image cache instead of rebuilding it (engine#1716) The int8 head tiles were gathered and packed on every serve. The file format (#1802) and the cache key (#1832) were merged but nothing called them. q4nx_head_tile_image() fills one head tile: a cached image is used when it reads back whole and a sample of its columns re-packs to the same bytes from the weights held now; anything else packs as before and writes the file. The forward's head loop calls it through I8Ctx::fill_weight_bo. Files live in /head-images; ONEBIT_HEAD_IMAGE_CACHE=0 turns it off. Fixes #1716 Co-Authored-By: Claude Fable 5.2 --- CMakeLists.txt | 9 +- docs/npu.md | 28 +++ npu/forward/engine-src/npu_engine_i8ctx_inc.h | 9 + npu/forward/include/q4nx_head_image.h | 48 ++++- npu/forward/src/q4nx_forward.cpp | 78 +++++++- npu/forward/src/q4nx_head_tile_cache.cpp | 170 ++++++++++++++++ tests/npu_design_grid_test.c | 186 ++++++++++++++++++ 7 files changed, 515 insertions(+), 13 deletions(-) create mode 100644 npu/forward/src/q4nx_head_tile_cache.cpp diff --git a/CMakeLists.txt b/CMakeLists.txt index e20a86864..19fc68b0f 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -1328,8 +1328,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 @@ -1569,6 +1573,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 diff --git a/docs/npu.md b/docs/npu.md index 6862549d4..19ebe3816 100644 --- a/docs/npu.md +++ b/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 `/head-images/` as +`headimg--t.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`. diff --git a/npu/forward/engine-src/npu_engine_i8ctx_inc.h b/npu/forward/engine-src/npu_engine_i8ctx_inc.h index 506d6cddf..356526609 100644 --- a/npu/forward/engine-src/npu_engine_i8ctx_inc.h +++ b/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 + 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) { diff --git a/npu/forward/include/q4nx_head_image.h b/npu/forward/include/q4nx_head_image.h index 573f48c82..98b7267bb 100644 --- a/npu/forward/include/q4nx_head_image.h +++ b/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 #include #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 `/headimg--t.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 diff --git a/npu/forward/src/q4nx_forward.cpp b/npu/forward/src/q4nx_forward.cpp index b587a1eb6..b5dd8f756 100644 --- a/npu/forward/src/q4nx_forward.cpp +++ b/npu/forward/src/q4nx_forward.cpp @@ -37,6 +37,7 @@ #include "residual-dump.h" #include #include +#include #include #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 @@ -5371,7 +5373,6 @@ bool Q4nxNpuForward::run_lm_head(const std::vector& h, std::vector if (head_packed && ((int)lm_head_packed_.size() != ntiles || lm_head_packed_T_ != T)) { lm_head_packed_.clear(); - std::vector 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: @@ -5398,20 +5399,76 @@ bool Q4nxNpuForward::run_lm_head(const std::vector& h, std::vector 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(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. @@ -5420,8 +5477,9 @@ bool Q4nxNpuForward::run_lm_head(const std::vector& h, std::vector 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 Bt; diff --git a/npu/forward/src/q4nx_head_tile_cache.cpp b/npu/forward/src/q4nx_head_tile_cache.cpp new file mode 100644 index 000000000..184d879d3 --- /dev/null +++ b/npu/forward/src/q4nx_head_tile_cache.cpp @@ -0,0 +1,170 @@ +// 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 lm-head tile cache (engine#1716): load one packed head tile from a file a previous serve wrote, or pack it and +// write that file. The contract is in q4nx_head_image.h; this file is the only place that decides "hit". +// +// It is separate from the forward so it runs without a device: the forward hands it the weight BO's mapping as `codes`, +// a test hands it a malloc'd buffer, and both get the same bytes. + +#include "q4nx_head_image.h" +#include "q4nx_pack.h" +#include "q4nx_weight_layout.h" + +#include +#include +#include + +namespace { + +inline size_t bpos(size_t i, size_t j, int KD, int ND, int blocked) { + return blocked ? q4nx_layout::blocked_index(i, j, KD, ND) : q4nx_layout::rowmajor_index(i, j, KD, ND); +} + +// The sampled columns of a tile holding `ncols` real columns: both ends and an even spread between them. +int sample_columns(int ncols, int out[Q4NX_HEAD_TILE_SAMPLE]) { + if (ncols <= 0) return 0; + int n = 0; + for (int s = 0; s < Q4NX_HEAD_TILE_SAMPLE; s++) { + const int j = Q4NX_HEAD_TILE_SAMPLE > 1 + ? (int)(((long long)(ncols - 1) * s) / (Q4NX_HEAD_TILE_SAMPLE - 1)) + : 0; + if (n == 0 || out[n - 1] != j) out[n++] = j; + } + return n; +} + +// Does the image agree with the packer on the sampled columns? Each column is packed ALONE through the production +// packer (K x 1, row-major) and compared with the image at the layout's own index, so nothing here restates the +// quantiser: if the packer changes, this changes with it. +bool sample_agrees(const Q4nxHeadImageKey& key, int n0, int ncols, int k, int kd, int nd, Q4nxHeadColFn col, void* ctx, + const int8_t* codes, const float* scales, Q4nxHeadTileResult* res) { + int cols[Q4NX_HEAD_TILE_SAMPLE]; + const int ns = sample_columns(ncols, cols); + std::vector w((size_t)k); + std::vector c((size_t)k); + for (int s = 0; s < ns; s++) { + const int j = cols[s]; + std::fill(w.begin(), w.end(), 0.0f); + col(ctx, n0 + j, w.data()); + float sc = 0.0f; + double ssum = 0.0; + if (q4nx_pack_B_int8_blocked(c.data(), w.data(), k, 1, k, 1, /*blocked=*/0, &sc, &ssum) != 0) return false; + if (std::memcmp(&sc, &scales[j], sizeof sc) != 0) { + std::snprintf(res->why, sizeof res->why, + "the cached head image's column %d has scale %.9g where these weights pack to %.9g: it was " + "written for different head weights or by a different packer", + n0 + j, (double)scales[j], (double)sc); + return false; + } + for (int i = 0; i < k; i++) { + const int8_t got = codes[bpos((size_t)i, (size_t)j, kd, nd, (int)key.blocked)]; + if (got != c[(size_t)i]) { + std::snprintf(res->why, sizeof res->why, + "the cached head image's column %d row %d holds code %d where these weights pack to %d: " + "it was written for different head weights or by a different packer", + n0 + j, i, (int)got, (int)c[(size_t)i]); + return false; + } + } + res->sampled++; + } + return true; +} + +} // namespace + +extern "C" uint64_t q4nx_head_source_hash(const void* data, size_t nbytes, const uint32_t* extra, int n_extra) { + if (!data || nbytes == 0) return 0; + uint64_t h = 1469598103934665603ull; + const unsigned char* b = (const unsigned char*)data; + for (size_t i = 0; i < nbytes; ++i) { + h ^= (uint64_t)b[i]; + h *= 1099511628211ull; + } + const unsigned char* e = (const unsigned char*)extra; + for (size_t i = 0; extra && i < (size_t)(n_extra > 0 ? n_extra : 0) * sizeof(uint32_t); ++i) { + h ^= (uint64_t)e[i]; + h *= 1099511628211ull; + } + return h ? h : 1; // 0 is "no data" +} + +extern "C" 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) { + Q4nxHeadTileResult local; + if (!res) res = &local; + std::memset(res, 0, sizeof *res); + if (!key || !col || !codes || !scales || tile < 0 || n0 < 0 || ncols <= 0 || k <= 0 || t <= 0 || ncols > t || k > kd || + t > nd || (uint32_t)k != key->kpad || (uint32_t)t != key->t) { + std::snprintf(res->why, sizeof res->why, "head tile %d: unusable geometry (k %d t %d kd %d nd %d ncols %d)", tile, k, + t, kd, nd, ncols); + return 0; + } + if (key->blocked && !q4nx_layout::blocked_fits(kd, nd)) { + std::snprintf(res->why, sizeof res->why, "head tile %d: %d x %d has no blocked layout", tile, kd, nd); + return 0; + } + const size_t image_bytes = (size_t)kd * (size_t)nd; + + char path[4096]; + path[0] = 0; + if (dir && *dir) { + char name[64]; + if (q4nx_head_image_key_name(key, name, (int)sizeof name)) { + // "headimg-.q4hi" -> "headimg--t.q4hi": one file per tile of the same key. + char* dot = std::strrchr(name, '.'); + if (dot) *dot = 0; + const int n = std::snprintf(path, sizeof path, "%s/%s-t%03d.q4hi", dir, name, tile); + if (n <= 0 || n >= (int)sizeof path) path[0] = 0; + } + } + + if (path[0]) { + Q4nxHeadImageInfo info; + char why[256]; + if (q4nx_head_image_read(path, codes, (int)image_bytes, scales, nd, &info, why, (int)sizeof why)) { + if (sample_agrees(*key, n0, ncols, k, kd, nd, col, ctx, codes, scales, res)) { + res->hit = 1; + return 1; + } + res->refused = 1; // `why` was set by the sample + } else if (!std::strstr(why, "not there")) { + res->refused = 1; + std::snprintf(res->why, sizeof res->why, "%s", why); + } + } + + // MISS: exactly what I8Ctx::packB_into() does -- zero the image, gather the tile, pack it. + std::memset(codes, 0, image_bytes); + std::fill(scales, scales + nd, 0.0f); + std::vector bt((size_t)k * (size_t)t, 0.0f); // padded columns must be zero + { + std::vector w((size_t)k); + for (int j = 0; j < ncols; j++) { + std::fill(w.begin(), w.end(), 0.0f); + col(ctx, n0 + j, w.data()); + for (int i = 0; i < k; i++) bt[(size_t)i * (size_t)t + (size_t)j] = w[(size_t)i]; + } + } + double ssum = 0.0; + if (q4nx_pack_B_int8_blocked(codes, bt.data(), k, t, kd, nd, (int)key->blocked, scales, &ssum) != 0) { + std::snprintf(res->why, sizeof res->why, "head tile %d: the packer refused %d x %d into %d x %d", tile, k, t, kd, nd); + return 0; + } + if (path[0]) res->wrote = q4nx_head_image_write(path, codes, scales, kd, nd); + return 1; +} diff --git a/tests/npu_design_grid_test.c b/tests/npu_design_grid_test.c index b3447d81f..92c9dc1be 100644 --- a/tests/npu_design_grid_test.c +++ b/tests/npu_design_grid_test.c @@ -604,8 +604,194 @@ static void test_head_image(void) { } } +/* ── THE HEAD TILE CACHE (engine#1716): what a serve calls instead of packing ───────────────────────────────────────── + * The issue is "the packed lm-head is rebuilt on every serve". The arms below are that sentence and its two dangerous + * neighbours, with a COUNTER on the gather so they are about work done and not about a log line: + * - a second pass over the same tiles gathers only the sample, and returns the bytes a fresh pack returns; + * - a changed key (layout, head bytes) packs again rather than borrowing the other image; + * - a damaged file, and an intact file written for OTHER weights, are both replaced rather than used. */ +typedef struct { + int nv, k; + int variant; /* 0: the fixture; 1: row 3 negated (same column scales, different codes); 2: everything x1.25 */ + long calls; /* columns gathered */ +} HeadFixture; + +static void head_fixture_col(void* ctx, int n, float* out) { + HeadFixture* f = (HeadFixture*)ctx; + f->calls++; + for (int k = 0; k < f->k; k++) { + float v = 0.0f; + if (n < f->nv) { + /* a distinct absmax per column (the lesson of #1778: equal column scales hide a per-column mistake) */ + v = ((float)((k * 37 + n * 11) % 211) - 105.0f) * 0.004f * (1.0f + 0.01f * (float)(n % 97)); + if (f->variant == 1 && k == 3) v = -v; + if (f->variant == 2) v *= 1.25f; + } + out[k] = v; + } +} + +/* What I8Ctx::packB_into() writes for one tile: zero the image, gather the tile, pack it. */ +static void head_reference(HeadFixture* f, int n0, int ncols, int k, int t, int kd, int nd, int blocked, int8_t* codes, + float* scales) { + float* bt = calloc((size_t)k * t, sizeof(float)); + float* col = malloc((size_t)k * sizeof(float)); + const long keep = f->calls; + for (int j = 0; j < ncols; j++) { + head_fixture_col(f, n0 + j, col); + for (int i = 0; i < k; i++) bt[(size_t)i * t + j] = col[i]; + } + f->calls = keep; + memset(codes, 0, (size_t)kd * nd); + for (int j = 0; j < nd; j++) scales[j] = 0.0f; + double ssum = 0.0; + q4nx_pack_B_int8_blocked(codes, bt, k, t, kd, nd, blocked, scales, &ssum); + free(bt); free(col); +} + +static void head_tile_path(const char* dir, const Q4nxHeadImageKey* key, int tile, char* out, size_t out_n) { + char name[64]; + q4nx_head_image_key_name(key, name, (int)sizeof name); + char* dot = strrchr(name, '.'); + if (dot) *dot = 0; + snprintf(out, out_n, "%s/%s-t%03d.q4hi", dir, name, tile); +} + +static void test_head_tile_cache(void) { + const int k = 128, t = 128, kd = 128, nd = 128, nv = 300, ntiles = 3; /* the last tile holds 44 real columns */ + char dir[] = "head_tile_cache.XXXXXX"; + check(mkdtemp(dir) != NULL, "head tile cache: a scratch cache directory"); + int8_t* codes = malloc((size_t)kd * nd); + int8_t* ref = malloc((size_t)kd * nd); + float* scales = malloc((size_t)nd * sizeof(float)); + float* sref = malloc((size_t)nd * sizeof(float)); + check(codes && ref && scales && sref, "head tile cache: the fixture allocates"); + if (!codes || !ref || !scales || !sref) return; + HeadFixture fx = {nv, k, 0, 0}; + char paths[8][512]; + int npaths = 0; + + for (int blocked = 0; blocked <= 1; blocked++) { + Q4nxHeadImageKey key = {(uint32_t)k, (uint32_t)t, (uint32_t)blocked, 0x1716c0ffeeull}; + int cold_ok = 0, cold_same = 0, warm_hit = 0, warm_same = 0; + long cold_calls = 0, warm_calls = 0; + for (int pass = 0; pass < 2; pass++) { + for (int tile = 0; tile < ntiles; tile++) { + const int n0 = tile * t, ncols = (nv - n0 < t) ? nv - n0 : t; + head_reference(&fx, n0, ncols, k, t, kd, nd, blocked, ref, sref); + Q4nxHeadTileResult r; + const long before = fx.calls; + memset(codes, 0x55, (size_t)kd * nd); /* the BO is not ours to assume zero on a hit */ + const int rc = q4nx_head_tile_image(dir, &key, tile, n0, ncols, k, t, kd, nd, head_fixture_col, &fx, codes, + scales, &r); + const int same = memcmp(codes, ref, (size_t)kd * nd) == 0 && memcmp(scales, sref, (size_t)t * sizeof(float)) == 0; + if (pass == 0) { + cold_ok += rc == 1 && r.hit == 0 && r.wrote == 1 && r.refused == 0; + cold_same += same; + cold_calls += fx.calls - before; + if (npaths < 8) head_tile_path(dir, &key, tile, paths[npaths++], sizeof paths[0]); + } else { + warm_hit += rc == 1 && r.hit == 1 && r.wrote == 0 && r.refused == 0 && r.sampled > 0; + warm_same += same; + warm_calls += fx.calls - before; + } + } + } + check_int(cold_ok, ntiles, blocked ? "head tile cache (b8): a first pass packs and writes every tile" + : "head tile cache (rm): a first pass packs and writes every tile"); + check_int(cold_same, ntiles, "... and its bytes are packB_into()'s bytes (codes and scales)"); + check_int((int)cold_calls, nv, "... having gathered every vocabulary column once"); + check_int(warm_hit, ntiles, "head tile cache: a SECOND pass takes every tile from the cache"); + check_int(warm_same, ntiles, "... byte-identical to a fresh pack, over a buffer that was not zero"); + check(warm_calls > 0 && warm_calls <= (long)ntiles * Q4NX_HEAD_TILE_SAMPLE, + "... gathering only the verification sample, not the vocabulary (the head is NOT rebuilt)"); + } + + /* a different key must not borrow the image: the layout and the head bytes each name a different file */ + { + Q4nxHeadImageKey other = {(uint32_t)k, (uint32_t)t, 0, 0x1716c0ffeeull ^ 1ull}; + Q4nxHeadTileResult r; + const long before = fx.calls; + const int rc = q4nx_head_tile_image(dir, &other, 0, 0, t, k, t, kd, nd, head_fixture_col, &fx, codes, scales, &r); + check(rc == 1 && r.hit == 0 && r.wrote == 1, "head tile cache: a different head hash packs, it does not hit"); + check_int((int)(fx.calls - before), t, "... gathering the whole tile"); + if (npaths < 8) head_tile_path(dir, &other, 0, paths[npaths++], sizeof paths[0]); + } + + Q4nxHeadImageKey key = {(uint32_t)k, (uint32_t)t, 1, 0x1716c0ffeeull}; + char path[512]; + head_tile_path(dir, &key, 0, path, sizeof path); + + /* (a) a damaged file is replaced, with the reason, and the next serve hits again */ + { + FILE* f = fopen(path, "r+b"); + check(f != NULL, "head tile cache: the cached file is where the key says it is"); + if (f) { fseek(f, 24 + 777, SEEK_SET); int c = fgetc(f); fseek(f, 24 + 777, SEEK_SET); fputc((c ^ 0x21) & 0xff, f); fclose(f); } + head_reference(&fx, 0, t, k, t, kd, nd, 1, ref, sref); + Q4nxHeadTileResult r; + int rc = q4nx_head_tile_image(dir, &key, 0, 0, t, k, t, kd, nd, head_fixture_col, &fx, codes, scales, &r); + check(rc == 1 && r.hit == 0 && r.refused == 1 && r.wrote == 1, "head tile cache: a damaged image is REPLACED, not used"); + check(strstr(r.why, "hash") != NULL, "... and the reason names the hash"); + check(memcmp(codes, ref, (size_t)kd * nd) == 0, "... and the tile served is the freshly packed one"); + rc = q4nx_head_tile_image(dir, &key, 0, 0, t, k, t, kd, nd, head_fixture_col, &fx, codes, scales, &r); + check(rc == 1 && r.hit == 1, "... after which the cache hits again"); + } + /* (b) THE ARM THE PAYLOAD HASH CANNOT PASS FOR US: an intact, self-consistent image written for OTHER weights under + * the same name. Variant 1 keeps every column scale and changes codes; variant 2 changes the scales. Both must be + * caught by the sample and replaced -- copying either into a weight BO would answer plausibly and wrongly. */ + for (int variant = 1; variant <= 2; variant++) { + fx.variant = variant; + head_reference(&fx, 0, t, k, t, kd, nd, 1, ref, sref); + Q4nxHeadTileResult r; + int rc = q4nx_head_tile_image(dir, &key, 0, 0, t, k, t, kd, nd, head_fixture_col, &fx, codes, scales, &r); + check(rc == 1 && r.hit == 0 && r.refused == 1, + variant == 1 ? "head tile cache: an intact image for OTHER weights (same scales, other codes) is not used" + : "head tile cache: an intact image for OTHER weights (other scales) is not used"); + check(strstr(r.why, variant == 1 ? "holds code" : "has scale") != NULL, "... and the reason names what disagreed"); + check(memcmp(codes, ref, (size_t)kd * nd) == 0 && memcmp(scales, sref, (size_t)t * sizeof(float)) == 0, + "... and the tile served is packed from the weights actually held"); + rc = q4nx_head_tile_image(dir, &key, 0, 0, t, k, t, kd, nd, head_fixture_col, &fx, codes, scales, &r); + check(rc == 1 && r.hit == 1, "... and the replacement is what the next serve hits"); + } + fx.variant = 0; + /* (c) the cache is never required: no directory, or one that cannot be written, packs and says it wrote nothing */ + { + head_reference(&fx, 256, 44, k, t, kd, nd, 1, ref, sref); + Q4nxHeadTileResult r; + int rc = q4nx_head_tile_image(NULL, &key, 2, 256, 44, k, t, kd, nd, head_fixture_col, &fx, codes, scales, &r); + check(rc == 1 && r.hit == 0 && r.wrote == 0 && memcmp(codes, ref, (size_t)kd * nd) == 0, + "head tile cache: with no cache directory the tile packs as it always did"); + rc = q4nx_head_tile_image("head_tile_cache.no-such-dir/x", &key, 2, 256, 44, k, t, kd, nd, head_fixture_col, &fx, + codes, scales, &r); + check(rc == 1 && r.hit == 0 && r.wrote == 0 && r.refused == 0 && memcmp(codes, ref, (size_t)kd * nd) == 0, + "head tile cache: an unwritable directory is not an error either"); + /* and a geometry the key does not describe is refused outright rather than packed under the wrong name */ + rc = q4nx_head_tile_image(dir, &key, 0, 0, t, k - 64, t, kd, nd, head_fixture_col, &fx, codes, scales, &r); + check(rc == 0 && r.why[0] != 0, "head tile cache: a K that is not the key's kpad is refused"); + } + /* (d) the source hash moves with the bytes and with the facts beside them */ + { + unsigned char bytes[64]; + for (int i = 0; i < 64; i++) bytes[i] = (unsigned char)(i * 7); + uint32_t extra[2] = {300, 128}; + const uint64_t h0 = q4nx_head_source_hash(bytes, sizeof bytes, extra, 2); + bytes[40] ^= 1; + const uint64_t h1 = q4nx_head_source_hash(bytes, sizeof bytes, extra, 2); + bytes[40] ^= 1; + extra[0] = 301; + const uint64_t h2 = q4nx_head_source_hash(bytes, sizeof bytes, extra, 2); + check(h0 != 0 && h1 != h0 && h2 != h0 && h1 != h2, "head source hash: one byte, or one fact beside the bytes, changes it"); + check(q4nx_head_source_hash(NULL, 0, NULL, 0) == 0, "head source hash: no data is 0, which names no cache entry"); + } + + for (int i = 0; i < npaths; i++) remove(paths[i]); + check(rmdir(dir) == 0, "head tile cache: only the files the keys name were written (the directory empties)"); + free(codes); free(ref); free(scales); free(sref); +} + int main(void) { test_head_image(); + test_head_tile_cache(); oga_geometry(); on_grid_geometry(); activation_pad();