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;