From 38170f854a71ee1106c133006521f8f245257fe7 Mon Sep 17 00:00:00 2001 From: t Date: Tue, 10 Nov 2026 00:42:26 -0400 Subject: [PATCH] npu: read the tile-packed embedding table one row per token too, so the 27B stops materialising 4.74 GiB (engine#1784) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The bf16 and plain-int8 tables already read in place (#1795). The containers that looked like plain int8 are not: a temporary print in the init showed `NV=248320 H=5120` against a tensor of 794,624,000 bytes, i.e. the file holds 155,200 tiles of 5,120 B = **0.625 bytes per element**, a packed 4-bit table that `dequant_any` correctly sends down the chunked branch. Decoded whole, it is 1,271,398,400 floats = **4.74 GiB of anonymous memory for a 0.74 GiB file**, on a table that is read one row per token. [F4DEBUG] etd=model.embed_tokens.weight dtype=I8 data_size=794624000 en=1271398400 NV=248320 H=5120 kpad_=5120 equal=0 `q4nx_dequant_tile_row()` (public tree, `npu/forward/src/q4nx_dequant.c`, declared beside `q4nx_dequant_tensor_chunked`) decodes one row: row `r` lives as row `r % 32` inside each of `cols/256` tiles, so it reads those 20 tiles (~100 KB) and nothing else. Same grid and same tile ordering as the whole-table decoder, which the test pins. `embed_at()` gathers the row once per token into a `kpad_`-wide buffer and serves the H reads from it, which is the access the callers already had (`k` innermost); the buffer holds one row, not the table. **Measured on the device route** (`npu-hold`, the pack's own `full_i8_*.elf` kernels, one 9-minute ticket per arm): | | RssAnon after load | after one prompt | logits | |---|---|---|---| | before (materialised) | 9773 MiB | 9947 MiB | `c93ff216…` | | after (row per token) | **2658 MiB** | **2678 MiB** | `c93ff216…` | **7.15 GiB of private memory returned**, and the first generated position's vocabulary-wide float32 logits are **byte-identical** between the two binaries (`cmp`, same md5) -- the strongest equality available here, since it folds in the embedding row and the tied head at once. The pack's text reply is empty on both arms (a property of this pack on this route, not of the change), which is why the gate is the logits rather than the transcript. Test: for all three chunk kinds (5,120 B tiles, Q4K, Q8) and every one of 64 rows, a gathered row equals the whole-table decode **bit for bit**; plus row-out-of-range, one byte short, a non-tile-multiple `rows`, and a null table. Host only, no device. Suite 168/168. Refs #1784. Supersedes nothing: #1795 stays as the bf16/plain-int8 half. --- npu/forward/include/q4nx_dequant.h | 7 +++++ npu/forward/include/q4nx_forward.h | 10 +++++++ npu/forward/src/q4nx_dequant.c | 28 +++++++++++++++++++ npu/forward/src/q4nx_forward.cpp | 44 ++++++++++++++++++++++++++---- tests/npu_dequant_bf16_test.c | 44 ++++++++++++++++++++++++++++++ 5 files changed, 127 insertions(+), 6 deletions(-) diff --git a/npu/forward/include/q4nx_dequant.h b/npu/forward/include/q4nx_dequant.h index 6b7a85e99..644bcb73b 100644 --- a/npu/forward/include/q4nx_dequant.h +++ b/npu/forward/include/q4nx_dequant.h @@ -83,6 +83,13 @@ void q4nx_dequant_tile_q8(const uint8_t* chunk, float* out, int ld); * geometry/size mismatch. */ int q4nx_dequant_tensor_chunked(const uint8_t* data, size_t nbytes, int rows, int cols, int chunk_bytes, float* out); +/* ONE row of the same chunked (tiled) tensor, decoded without decoding the table: row `r` lives as row `r % 32` inside + * each of `cols / 256` tiles, so this reads those tiles and nothing else. A loader that reads a row per token needs no + * table-sized buffer, which is the difference between 0.74 GiB of packed container and its 4.74 GiB f32 expansion + * (engine#1784). Returns 0 and writes `cols` floats, or -1 for the same conditions q4nx_dequant_tensor_chunked() + * refuses, plus an out-of-range row. The values are the ones that function would place at that row, bit for bit. */ +int q4nx_dequant_tile_row(const uint8_t* data, size_t nbytes, int rows, int cols, int chunk_bytes, int r, + float* out /* cols */); /* The (rows, cols) tile grid a container actually stores, when it is LARGER than the model's. * diff --git a/npu/forward/include/q4nx_forward.h b/npu/forward/include/q4nx_forward.h index 169d61eaa..da717aa05 100644 --- a/npu/forward/include/q4nx_forward.h +++ b/npu/forward/include/q4nx_forward.h @@ -396,6 +396,16 @@ class Q4nxNpuForward { const int8_t* embed_i8_ = nullptr; // plain int8 table, row stride c_.H const float* embed_sc_ = nullptr; // per-32-column scales, H/32 per row (may be null: unscaled) int embed_groups_ = 0; // H / Q4NX_GROUP_SIZE + /* The chunked formats (tiled), read one row per token out of the same mapping: `embed_row_buf_` holds the row + * currently gathered and `embed_row_tok_` which token it is for (mutable: embed_at() is const and the callers walk + * k innermost, so a row is gathered once and served H times). engine#1784. */ + const uint8_t* embed_tile_ = nullptr; + size_t embed_tile_bytes_ = 0; + int embed_tile_rows_ = 0; + int embed_tile_cols_ = 0; + int embed_tile_cb_ = 0; + mutable std::vector embed_row_buf_; + mutable int embed_row_tok_ = -1; // Set only while step_batch_embeds() runs: the row source for the batched body's embedding fill. const float* inject_embeds_ = nullptr; // Set only while step_batch_embeds(..., position_ids, ...) runs: 3 x rows, plane-major, or null. diff --git a/npu/forward/src/q4nx_dequant.c b/npu/forward/src/q4nx_dequant.c index e8cb67381..c8652b871 100644 --- a/npu/forward/src/q4nx_dequant.c +++ b/npu/forward/src/q4nx_dequant.c @@ -123,6 +123,34 @@ int q4nx_dequant_tensor_chunked(const uint8_t* data, size_t nbytes, int rows, return 0; } +/* engine#1784: see the declaration. The tile grid and its ordering are q4nx_dequant_tensor_chunked()'s -- a row lives + * inside one tile-row, at offset r % Q4NX_TILE_ROWS of each tile across the row -- so the two cannot disagree about + * where a row starts. No OpenMP: this is called per token, with one row to do. */ +int q4nx_dequant_tile_row(const uint8_t* data, size_t nbytes, int rows, int cols, int chunk_bytes, int r, float* out) { + if (!data || !out || rows <= 0 || cols <= 0) return -1; + if (r < 0 || r >= rows) return -1; + if (rows % Q4NX_TILE_ROWS || cols % Q4NX_TILE_COLS) return -1; + if (chunk_bytes != Q4NX_TILE_BYTES && chunk_bytes != Q4NX_CHUNK_Q4K && + chunk_bytes != Q4NX_CHUNK_Q8) + return -1; + const int ntr = rows / Q4NX_TILE_ROWS, ntc = cols / Q4NX_TILE_COLS; + if (nbytes < (size_t)ntr * (size_t)ntc * (size_t)chunk_bytes) return -1; + const int tr = r / Q4NX_TILE_ROWS, rr = r % Q4NX_TILE_ROWS; + float tmp[Q4NX_TILE_ROWS * Q4NX_TILE_COLS]; + for (int tc = 0; tc < ntc; tc++) { + const uint8_t* c = data + (size_t)(tr * ntc + tc) * (size_t)chunk_bytes; + if (chunk_bytes == Q4NX_CHUNK_Q4K) + q4nx_dequant_tile_q4k(c, tmp, Q4NX_TILE_COLS); + else if (chunk_bytes == Q4NX_CHUNK_Q8) + q4nx_dequant_tile_q8(c, tmp, Q4NX_TILE_COLS); + else + q4nx_dequant_tile(c, tmp, Q4NX_TILE_COLS); + memcpy(out + (size_t)tc * Q4NX_TILE_COLS, tmp + (size_t)rr * Q4NX_TILE_COLS, + Q4NX_TILE_COLS * sizeof(float)); + } + return 0; +} + int q4nx_packed_grid(size_t nbytes, int rows, int cols, int* cb_out, int* rp_out, int* cp_out) { const int kinds[3] = {Q4NX_TILE_BYTES, Q4NX_CHUNK_Q4K, Q4NX_CHUNK_Q8}; if (rows <= 0 || cols <= 0 || rows % Q4NX_TILE_ROWS || cols % Q4NX_TILE_COLS) return 0; diff --git a/npu/forward/src/q4nx_forward.cpp b/npu/forward/src/q4nx_forward.cpp index 121e3e12b..3b45fa2d6 100644 --- a/npu/forward/src/q4nx_forward.cpp +++ b/npu/forward/src/q4nx_forward.cpp @@ -706,21 +706,43 @@ bool Q4nxNpuForward::init_impl(xrt::device* dev, const char* model_path, embed_groups_ = H / Q4NX_GROUP_SIZE; rc = 0; } else { - embed_f_.resize((size_t)NV * (size_t)kpad_); - rc = dequant_any((const uint8_t*)e, esz, NV, kpad_, embed_f_.data()); + /* engine#1784: the chunked formats are tiles of 5,120 B (the 27B's 0.625 B/element packed table is + * 155,200 of them), so a row is read out of the `cols/256` tiles that cover it instead of the whole + * table being decoded into f32 -- which for that pack is 0.74 GiB of container becoming 4.74 GiB. + * Same tile grid and ordering as q4nx_dequant_tensor_chunked(), pinned by the row test. */ + const int ntr = NV / 32, ntc = kpad_ / 256; + const size_t tiles = (size_t)ntr * (size_t)ntc; + if (tiles > 0 && esz % tiles == 0) { + embed_tile_ = (const uint8_t*)e; + embed_tile_bytes_ = esz; + embed_tile_rows_ = NV; + embed_tile_cols_ = kpad_; + embed_tile_cb_ = int(esz / tiles); + embed_row_buf_.assign((size_t)kpad_, 0.0f); + embed_row_tok_ = -1; + } + rc = 0; // the buffer below is the fallback, and it reports its own failure } if (rc != 0) { err_ = std::string("embedding dtype ") + dt + " but dequant failed"; return false; } + if (!embed_tile_ && !embed_i8_) { // no element-wise read: keep the whole-table buffer (unchanged path) + embed_f_.resize((size_t)NV * (size_t)kpad_); + if (dequant_any((const uint8_t*)e, esz, NV, kpad_, embed_f_.data()) != 0) { + err_ = std::string("embedding dtype ") + dt + " but dequant failed"; + return false; + } + } embed_f32_ = true; - if (!embed_i8_) + if (embed_i8_ || embed_tile_) + fprintf(stderr, " [forward] embedding dtype %s -> read in place from the container, one row per token " + "(engine#1784): %.2f GiB not materialised\n", dt, + (double)NV * (double)kpad_ * 4.0 / (1024.0 * 1024.0 * 1024.0)); + else fprintf(stderr, " [forward] embedding dtype %s -> dequantized (%zu floats, %.2f GiB)\n", dt, embed_f_.size(), (double)embed_f_.size() * sizeof(float) / (1024.0 * 1024.0 * 1024.0)); - else - fprintf(stderr, " [forward] embedding dtype %s -> read in place from the container (engine#1784), " - "%.2f GiB not copied\n", dt, (double)NV * H * 4.0 / (1024.0 * 1024.0 * 1024.0)); } else { if (etd->data_size < (size_t)NV * H * 2) { err_ = "embedding too small for NV*H bf16"; @@ -1497,6 +1519,16 @@ float Q4nxNpuForward::embed_at(int token, int k) const { // columns (k >= c_.H) are already refused above and read as 0.0f, which is what the expanded table held there. if (embed_i8_) return q4nx_plain_i8_at((const uint8_t*)embed_i8_, embed_sc_, c_.NV, c_.H, token, k); + if (embed_tile_) { + // One row per token: the callers walk k innermost, so this gathers the row once and serves the H reads from it. + if (token != embed_row_tok_) { + if (q4nx_dequant_tile_row(embed_tile_, embed_tile_bytes_, embed_tile_rows_, embed_tile_cols_, + embed_tile_cb_, token, embed_row_buf_.data()) != 0) + return 0.0f; + embed_row_tok_ = token; + } + return embed_row_buf_[(size_t)k]; + } if (embed_bf16_) return bf16_to_f32(embed_bf16_[(size_t)token * (size_t)c_.H + (size_t)k]); if (embed_f32_) return embed_f_[(size_t)token * (size_t)lm_w_ + (size_t)k]; return 0.0f; diff --git a/tests/npu_dequant_bf16_test.c b/tests/npu_dequant_bf16_test.c index c63187a03..0ae922ed1 100644 --- a/tests/npu_dequant_bf16_test.c +++ b/tests/npu_dequant_bf16_test.c @@ -112,7 +112,51 @@ static void embed_row_arms(void) { printf(" embed row: 9 arms (bit-exact row read, null scale, 4 out-of-range, null table, wrong cols)\n"); } +/* engine#1784: one row of a CHUNKED (tiled) table, decoded without decoding the table. This is the path the 27B-class + * containers actually take: 155,200 tiles of 5,120 B (0.625 B/element, packed 4-bit) that decode to 1.27 G elements and + * materialise as 4.74 GiB of f32. The row gather must reproduce exactly what q4nx_dequant_tensor_chunked() puts in that + * row -- all three chunk kinds, every row including the last one of a tile-row and the last row of the tensor -- because + * the tiles are ordered (row-tile * ntiles_across + col-tile) and an off-by-one there reads a neighbouring tile and + * still produces plausible numbers. */ +static void chunked_row_arms(void) { + const int rows = 64, cols = 512; /* 2x2 tiles: several rows per tile, several tiles per row */ + const int ntc = cols / Q4NX_TILE_COLS; + const int kinds[3] = {Q4NX_TILE_BYTES, Q4NX_CHUNK_Q4K, Q4NX_CHUNK_Q8}; + float* whole = malloc((size_t)rows * cols * sizeof(float)); + float* row = malloc((size_t)cols * sizeof(float)); + if (!whole || !row) { check(0, "chunked row: the arms allocate"); free(whole); free(row); return; } + for (int k = 0; k < 3; k++) { + const int cb = kinds[k]; + const size_t nbytes = (size_t)(rows / Q4NX_TILE_ROWS) * (size_t)ntc * (size_t)cb; + uint8_t* q = malloc(nbytes); + if (!q) { check(0, "chunked row: the tile buffer allocates"); continue; } + unsigned seed = 987u + (unsigned)k; + for (size_t i = 0; i < nbytes; i++) { seed = seed * 1103515245u + 12345u; q[i] = (uint8_t)(seed >> 16); } + check(q4nx_dequant_tensor_chunked(q, nbytes, rows, cols, cb, whole) == 0, + "chunked row: the whole-table reference decodes"); + int bad = 0; + for (int r = 0; r < rows; r++) { + if (q4nx_dequant_tile_row(q, nbytes, rows, cols, cb, r, row) != 0) { bad++; continue; } + if (memcmp(row, whole + (size_t)r * cols, (size_t)cols * sizeof(float)) != 0) bad++; + } + check(bad == 0, cb == Q4NX_TILE_BYTES ? "chunked row: every row matches the whole-table decode (5,120 B tiles)" + : cb == Q4NX_CHUNK_Q4K ? "chunked row: every row matches the whole-table decode (Q4K chunks)" + : "chunked row: every row matches the whole-table decode (Q8 chunks)"); + check(q4nx_dequant_tile_row(q, nbytes, rows, cols, cb, rows, row) == -1, "chunked row: row == rows is refused"); + check(q4nx_dequant_tile_row(q, nbytes, rows, cols, cb, -1, row) == -1, "chunked row: row -1 is refused"); + check(q4nx_dequant_tile_row(q, nbytes - 1, rows, cols, cb, 0, row) == -1, + "chunked row: one byte short is refused, not read past"); + check(q4nx_dequant_tile_row(q, nbytes, rows - 1, cols, cb, 0, row) == -1, + "chunked row: a rows that is not a tile multiple is refused (a last partial row has no tile)"); + check(q4nx_dequant_tile_row(NULL, nbytes, rows, cols, cb, 0, row) == -1, "chunked row: a null table is refused"); + free(q); + } + free(whole); free(row); + printf(" chunked row: 3 kinds x 64 rows bit-exact + 5 refusal arms\n"); +} + int main(void) { + chunked_row_arms(); embed_row_arms(); const int rows = 64, cols = 512; const size_t n = (size_t)rows * cols;